deepmd.dpmodel.array_api#
Utilities for the array API.
Attributes#
Functions#
|
|
| |
| |
| Contract operands according to the Einstein summation subscripts. |
| Array-API-only |
| Take the first n elements along dim. |
| Reduces all values from the src tensor to the indices specified in the index tensor. |
| Adds values to the specified indices of x in place or returns new x (for JAX). |
| Mark a data-dependent leading dimension as a valid size for torch.export. |
| Segment max-assign of values into x at the specified indices. |
| Compute the sigmoid function. |
| Compute the error function. |
| Set items at boolean mask indices. |
| Draw |
| Counts the number of occurrences of each value in x. |
Module Contents#
- deepmd.dpmodel.array_api.xp_asarray_nodetach(xp: Any, obj: Any, *, dtype: Any = None, device: Any = None) Array[source]#
xp.asarraythat preserves autograd for backend tensors.torch.asarraydetaches its input from the autograd graph, so callingxp.asarrayon a weight attribute that is already a backend tensor (e.g. atorch.nn.Parameterregistered by the pt_expt backend) silently breaks gradient flow to that weight. Backend tensors already inxpare therefore returned as-is, with an optional differentiable dtype cast viaxp.astype.An array from another namespace cannot retain its autograd graph. It is converted through NumPy before entering
xp; this also performs the required device-to-host copy when a CUDA-backed model constant is consumed by a NumPy statistics path.An array already in
xpnormally needs no move, since model buffers and inputs travel together. Tracing breaks that: a CPU export of a CUDA-resident model runs CPU inputs through a module whose buffers stayed on the device, and the mismatch surfaces far downstream as a fake-tensor device error. A requesteddevicethat differs is therefore honoured here, which costs a comparison on the hot path and a copy only in the tracing case.
- deepmd.dpmodel.array_api.xp_einsum(subscripts: str, *operands: Array) Array[source]#
Contract operands according to the Einstein summation subscripts.
The array API standard has no
einsum, so an array-API-only module has to express a contraction as a chain ofpermute_dims/reshape/matmul. That chain fixes one particular execution order, which is rarely the best one and is opaque to a compiler: a batched contraction written this way reaches PyTorch asbmmplus transposing copies, where the same expression given totorch.einsumis free to become a singlemmon a reshaped operand or to fuse into a neighbouring kernel. On the production DPA4 shapes the difference is measurable – the array-API chain put roughly 150 more contractions per training step onbmminstead ofmm, for about 5% of the step and an extra gigabyte of peak memory.Writing the chain by hand is also a correctness-adjacent hazard, because the orders differ by more than a constant. The broadcast spelling
matmul(x[..., None, :], weight[None, ...])makes the node count the matmul batch, so the weight is expanded across it and autograd reduces that whole expansion back to the parameter shape: at production sizes a 165 K-element weight became 191 M elements, and its reduce was the single costliest kernel of a training step (a 15x penalty on the affected contraction). Stating the contraction leaves that choice to the backend.Every backend this project targets (NumPy, PyTorch, JAX) ships an
einsum, so it is dispatched directly where available. The array-API fallback below serves the remaining namespaces (array_api_strict, used by the conformance tests) and is restricted to what those need.- Parameters:
- Returns:
ArrayThe contraction result.
- Raises:
ValueErrorIf subscripts is in implicit form, or if the fallback is reached with a specification it does not implement.
- deepmd.dpmodel.array_api._xp_einsum_fallback(subscripts: str, *operands: Array) Array[source]#
Array-API-only
einsumfor a two-operand contraction.Serves the namespaces without a native
einsum. The contraction is reduced to the canonical batched matmul: labels shared by both operands and the output are the batch, labels shared by the operands but absent from the output are contracted, and the rest are free on one side each. Each operand is permuted into(batch, free, contracted)order, flattened to three axes, multiplied, and restored to the requested output order.Correctness rather than throughput is the aim here: the flattening materializes a copy of each operand whenever the permutation is not a view, which a native
einsumwould avoid. That trade is deliberate – every backend used for production has aneinsum, and this path exists for the conformance namespaces.Only a diagonal (a label repeated within one operand) and an implicit output are rejected; both are absent from this codebase.
- deepmd.dpmodel.array_api.xp_take_first_n(arr: Array, dim: int, n: int) Array[source]#
Take the first n elements along dim.
For torch tensors, uses
torch.index_selectso thattorch.exportdoes not emit a contiguity guard that would prevent thenall == nloc(no-PBC) case from working. For numpy / jax, uses regular slicing.
- deepmd.dpmodel.array_api.xp_scatter_sum(input: Array, dim: int, index: Array, src: Array) Array[source]#
Reduces all values from the src tensor to the indices specified in the index tensor.
This function is similar to PyTorch’s scatter_add and JAX’s scatter_sum. It adds values from src to input at positions specified by index along the given dimension.
- deepmd.dpmodel.array_api.xp_add_at(x: Array, indices: Array, values: Array) Array[source]#
Adds values to the specified indices of x in place or returns new x (for JAX).
- deepmd.dpmodel.array_api.xp_hint_dynamic_size(x: Array) None[source]#
Mark a data-dependent leading dimension as a valid size for torch.export.
Under symbolic tracing (
make_fx/torch.export) the length of a data-dependent array (e.g. the output ofnonzeroor a tensor-repeat) is an UNBACKED SymInt; guarding Python control flow or allocations on it raisesGuardOnDataDependentSymNode.torch._check_is_sizeregisters the>= 0size hint that lets the tracer treat it as a proper dimension (recorded as asym_constrain_range_for_sizenode, preserved by AOTI).No-op for numpy / jax / eager-torch concrete shapes — safe to call unconditionally from dpmodel code (torch imported lazily, torch arrays only).
- deepmd.dpmodel.array_api.xp_maximum_at(x: Array, indices: Array, values: Array) Array[source]#
Segment max-assign of values into x at the specified indices.
Element-wise analogue of
xp_add_at()that takes the maximum instead of the sum: for everykit assignsx[indices[k]] = maximum( x[indices[k]], values[k]). Repeated indices reduce to the per-segment maximum, which is order-independent.- Parameters:
- Returns:
ArrayThe updated array (modified in place and returned for NumPy; a new array for JAX/PyTorch).
- deepmd.dpmodel.array_api.xp_sigmoid(x: Array) Array[source]#
Compute the sigmoid function.
JAX and PyTorch have optimized sigmoid implementations. See jax-ml/jax#15617
- deepmd.dpmodel.array_api.xp_erf(x: Array) Array[source]#
Compute the error function.
Used by the exact (non-approximated) GELU. The array API has no
erf, so each backend’s own implementation is used; NumPy goes through SciPy, which is already a core dependency.
- deepmd.dpmodel.array_api.xp_setitem_at(x: Array, mask: Array, values: Array) Array[source]#
Set items at boolean mask indices.
For JAX and PyTorch arrays, returns a new array (non-mutating). For NumPy arrays, modifies in-place and returns the same array.
- deepmd.dpmodel.array_api.xp_uniform(like: Array, size: int, low: float = 0.0, high: float = 1.0) Array[source]#
Draw
sizeuniform samples in[low, high)onlike’s device.Each backend uses its own generator: TensorFlow draws with
tf.random.uniformso traced graphs advance the runtime RNG, torch draws withtorch.rand(sosetup_seedreplays it without a host copy), and other backends usedeepmd.utils.random. Draws are therefore not comparable across backends – use only for a per-forward random stream, never where a parity test looks.