onnx-diagnostic 0.8.0__py3-none-any.whl

onnx_diagnostic/reference/__init__.py

@@ -0,0 +1,4 @@
+from .evaluator import ExtendedReferenceEvaluator
+from .ort_evaluator import OnnxruntimeEvaluator
+from .torch_evaluator import TorchOnnxEvaluator
+from .report_results_comparison import ReportResultComparison

onnx_diagnostic/reference/evaluator.py

@@ -0,0 +1,254 @@
+from logging import getLogger
+from typing import Any, Dict, List, Optional, Union
+from onnx import FunctionProto, ModelProto, NodeProto, TypeProto
+from onnx.defs import get_schema
+from onnx.reference import ReferenceEvaluator
+from onnx.reference.op_run import OpRun
+from .ops.op_add_add_mul_mul import (
+    AddAdd,
+    AddMul,
+    AddSharedInput,
+    MulAdd,
+    MulMul,
+    MulSharedInput,
+    MulSub,
+    SubMul,
+)
+from .ops.op_attention import Attention
+from .ops.op_average_pool_grad import AveragePoolGrad
+from .ops.op_bias_softmax import BiasSoftmax
+from .ops.op_cast_like import CastLike_15, CastLike_19
+from .ops.op_complex import ComplexModule, ToComplex
+from .ops.op_concat import Concat
+from .ops.op_constant_of_shape import ConstantOfShape
+from .ops.op_fused_matmul import FusedMatMul
+from .ops.op_gather import Gather
+from .ops.op_gather_elements import GatherElements
+from .ops.op_gather_grad import GatherGrad
+from .ops.op_memcpy_host import MemcpyFromHost, MemcpyToHost
+from .ops.op_mul_sigmoid import MulSigmoid
+from .ops.op_negxplus1 import NegXplus1
+from .ops.op_qlinear_average_pool import QLinearAveragePool
+from .ops.op_qlinear_conv import QLinearConv
+from .ops.op_quick_gelu import QuickGelu
+from .ops.op_replace_zero import ReplaceZero
+from .ops.op_rotary import Rotary
+from .ops.op_scan import Scan
+from .ops.op_scatter_elements import ScatterElements
+from .ops.op_scatternd_of_shape import MaskedScatterNDOfShape, ScatterNDOfShape
+from .ops.op_simplified_layer_normalization import SimplifiedLayerNormalization
+from .ops.op_skip_layer_normalization import SkipLayerNormalization
+from .ops.op_slice import Slice_1, Slice_10
+from .ops.op_transpose_cast import Transpose2DCastFP16, Transpose2DCastFP32
+from .ops.op_tri_matrix import TriMatrix
+
+
+logger = getLogger("onnx-diagnostic-eval")
+
+
+class ExtendedReferenceEvaluator(ReferenceEvaluator):
+    """
+    This class replaces the python implementation by custom implementation.
+    The evaluator allows to test
+    scenarios outside what an onnx backend bound to the official onnx
+    operators definition could do such as optimization patterns
+    involving onnxruntime contrib operators.
+
+    ::
+
+        from onnx_diagnostic.reference import ExtendedReferenceEvaluator
+        ref = ExtendedReferenceEvaluator(...)
+
+    The class overloads or adds the following operators by default:
+
+    .. runpython::
+        :showcode:
+
+        import pprint
+        from onnx_diagnostic.reference import ExtendedReferenceEvaluator
+
+        pprint.pprint(ExtendedReferenceEvaluator.default_ops)
+    """
+
+    default_ops: List[type[OpRun]] = [
+        AddAdd,
+        AddMul,
+        AddSharedInput,
+        Attention,
+        AveragePoolGrad,
+        BiasSoftmax,
+        Concat,
+        CastLike_15,
+        CastLike_19,
+        ComplexModule,
+        ConstantOfShape,
+        FusedMatMul,
+        Gather,
+        GatherElements,
+        GatherGrad,
+        MaskedScatterNDOfShape,
+        MemcpyFromHost,
+        MemcpyToHost,
+        MulAdd,
+        MulMul,
+        MulSharedInput,
+        MulSigmoid,
+        MulSub,
+        NegXplus1,
+        QLinearConv,
+        QLinearAveragePool,
+        QuickGelu,
+        ReplaceZero,
+        Rotary,
+        Scan,
+        ScatterElements,
+        ScatterNDOfShape,
+        SimplifiedLayerNormalization,
+        SkipLayerNormalization,
+        Slice_1,
+        Slice_10,
+        SubMul,
+        ToComplex,
+        Transpose2DCastFP16,
+        Transpose2DCastFP32,
+        TriMatrix,
+    ]
+
+    @staticmethod
+    def filter_ops(proto, new_ops, opsets):
+        if opsets is None and isinstance(proto, (ModelProto, FunctionProto)):
+            opsets = {d.domain: d.version for d in proto.opset_import}
+        best = {}
+        renamed = {}
+        for cl in new_ops:
+            if "_" not in cl.__name__:
+                continue
+            vers = cl.__name__.split("_")
+            try:
+                v = int(vers[-1])
+            except ValueError:
+                # not a version
+                continue
+            if opsets is not None and v > opsets.get(cl.op_domain, 1):
+                continue
+            renamed[cl.__name__] = cl
+            key = cl.op_domain, "_".join(vers[:-1])
+            if key not in best or best[key][0] < v:
+                best[key] = (v, cl)
+
+        modified = []
+        for cl in new_ops:
+            if cl.__name__ not in renamed:
+                modified.append(cl)
+        for k, v in best.items():
+            atts = {"domain": k[0]}
+            bases = (v[1],)
+            if not hasattr(v[1], "op_schema"):
+                atts["op_schema"] = get_schema(k[1], v[0], domain=v[1].op_domain)
+            new_cl = type(k[1], bases, atts)
+            modified.append(new_cl)
+
+        new_ops = modified
+        return new_ops
+
+    def __init__(
+        self,
+        proto: Any,
+        opsets: Optional[Dict[str, int]] = None,
+        functions: Optional[List[Union[ReferenceEvaluator, FunctionProto]]] = None,
+        verbose: int = 0,
+        new_ops: Optional[List[type[OpRun]]] = None,
+        **kwargs,
+    ):
+        if new_ops is None:
+            new_ops = ExtendedReferenceEvaluator.default_ops
+        else:
+            new_ops = new_ops.copy()
+            new_ops.extend(ExtendedReferenceEvaluator.default_ops)
+        new_ops = ExtendedReferenceEvaluator.filter_ops(proto, new_ops, opsets)
+
+        ReferenceEvaluator.__init__(
+            self,
+            proto,
+            opsets=opsets,
+            functions=functions,
+            verbose=verbose,
+            new_ops=new_ops,
+            **kwargs,
+        )
+
+    def _log(self, level: int, pattern: str, *args: List[Any]) -> None:
+        if level < self.verbose:
+            new_args = [self._log_arg(a) for a in args]
+            print(pattern % tuple(new_args))
+        else:
+            logger.debug(pattern, *args)
+
+    def run(self, *args, **kwargs):
+        """See :meth:`onnx.reference.ReferenceEvaluator.run`."""
+        if len(args) == 1 and isinstance(args[0], list):
+            feeds = dict(zip(self.input_names, args[0]))
+            return self.run(None, feeds, **kwargs)
+        if isinstance(self.proto_, FunctionProto):
+            return self._run_function(*args, **kwargs)
+        return ReferenceEvaluator.run(self, *args, **kwargs)
+
+    def _load_impl(self, node: NodeProto, input_types: TypeProto | None = None) -> Any:
+        res = super()._load_impl(node, input_types)
+        assert (
+            not hasattr(res, "op_domain") or res.op_domain == node.domain
+        ), f"Domain mismatch {res.op_domain!r} != {node.domain} for node={node}"
+        return res
+
+    def _run_function(
+        self,
+        output_names,
+        feed_inputs: Dict[str, Any],
+        attributes: Optional[Dict[str, Any]] = None,
+        intermediate: bool = False,
+    ) -> Union[Dict[str, Any], List[Any]]:  # type: ignore
+        if output_names is None:
+            output_names = self.output_names
+
+        # step 1: inputs and initializers
+        results = {"": None}  # optional input
+        results.update(self.rt_inits_)  # type: ignore[arg-type]
+        results.update(feed_inputs)
+        for k, v in self.rt_inits_.items():
+            self._log(2, " +C %s: %s", k, v)  # type: ignore[arg-type]
+        for k, v in feed_inputs.items():
+            self._log(2, " +I %s: %s", k, v)  # type: ignore[arg-type]
+
+        # step 2: execute nodes
+        for node in self.rt_nodes_:
+            self._log(1, "%s(%s) -> %s", node.op_type, node.input, node.output)
+            for i in node.input:
+                if i not in results:
+                    raise RuntimeError(
+                        f"Unable to find input {i!r} in known results {sorted(results)}, "
+                        f"self.rt_inits_ has {sorted(self.rt_inits_)}, "
+                        f"feed_inputs has {sorted(feed_inputs)}."
+                    )
+            inputs = [results[i] for i in node.input]
+            linked_attributes = {}
+            if node.has_linked_attribute and attributes:
+                linked_attributes["linked_attributes"] = attributes
+            if node.need_context():
+                outputs = node.run(*inputs, context=results, **linked_attributes)
+            else:
+                outputs = node.run(*inputs, **linked_attributes)
+            for name, value in zip(node.output, outputs):
+                self._log(2, " + %s: %s", name, value)  # type: ignore[arg-type]
+                results[name] = value
+
+        # return the results
+        if intermediate:
+            return results
+
+        for name in output_names:
+            if name not in results:
+                raise RuntimeError(
+                    f"Unable to find output name {name!r} "
+                    f"in {sorted(results)}, proto is\n{self.proto_}"
+                )
+        return [results[name] for name in output_names]

onnx_diagnostic/reference/ops/__init__.py

@@ -0,0 +1 @@
+

onnx_diagnostic/reference/ops/op_add_add_mul_mul.py

@@ -0,0 +1,68 @@
+import numpy as np
+from onnx.reference.op_run import OpRun
+
+
+class AddAdd(OpRun):
+    op_domain = "onnx_extended.ortops.optim.cuda"
+
+    def _run(self, x, y, z):
+        return (x + y + z,)
+
+
+class MulMul(OpRun):
+    op_domain = "onnx_extended.ortops.optim.cuda"
+
+    def _run(self, x, y, z):
+        return (x * y * z,)
+
+
+class AddMul(OpRun):
+    op_domain = "onnx_extended.ortops.optim.cuda"
+
+    def _run(self, x, y, z, transposeMiddle=None):
+        res = (x + y) * z
+        if transposeMiddle:
+            res = np.transpose(res, axes=[0, 2, 1, 3])
+        return (res,)
+
+
+class MulAdd(OpRun):
+    op_domain = "onnx_extended.ortops.optim.cuda"
+
+    def _run(self, x, y, z, transposeMiddle=None):
+        res = (x * y) + z
+        if transposeMiddle:
+            res = np.transpose(res, axes=[0, 2, 1, 3])
+        return (res,)
+
+
+class SubMul(OpRun):
+    op_domain = "onnx_extended.ortops.optim.cuda"
+
+    def _run(self, x, y, z, negative=None):
+        if negative:
+            return ((y - x) * z,)
+        return ((x - y) * z,)
+
+
+class MulSub(OpRun):
+    op_domain = "onnx_extended.ortops.optim.cuda"
+
+    def _run(self, x, y, z, negative=None):
+        if negative:
+            return (z - (x * y),)
+        return ((x * y) - z,)
+
+
+class AddSharedInput(OpRun):
+    op_domain = "onnx_extended.ortops.optim.cuda"
+
+    def _run(self, x, y, z):
+        return (x + y, x + z)
+
+
+class MulSharedInput(OpRun):
+    op_domain = "onnx_extended.ortops.optim.cuda"
+
+    def _run(self, x, y, z):
+        return (x * y, x * z)

onnx_diagnostic/reference/ops/op_attention.py

@@ -0,0 +1,60 @@
+import numpy as np
+import scipy.special as scipy_special
+from onnx.reference.op_run import OpRun
+
+
+class Attention(OpRun):
+    op_domain = "com.microsoft"
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        # Fix a bug onnx.reference.ReferenceEvaluator
+        self._schema = None
+        self.attributes_names_ = ["num_heads"]
+
+    def _run(
+        self,
+        x,
+        weights,
+        bias,
+        mask_index,
+        past,
+        attention_bias,
+        num_heads=None,
+    ):
+        assert past is None, f"Attention not implemented if past == {past!r}"
+        assert (
+            num_heads == attention_bias.shape[1]
+        ), f"num_heads={num_heads} not in attention_bias.shape={attention_bias.shape}"
+        d = weights.shape[1] // 3
+        q_weights = weights[:, :d]
+        k_weights = weights[:, d : d * 2]
+        v_weights = weights[:, d * 2 :]
+
+        d = bias.shape[0] // 3
+        q_bias = bias[:d]
+        k_bias = bias[d : d * 2]
+        v_bias = bias[d * 2 :]
+
+        shape_4d = (*x.shape[:2], num_heads, -1)
+
+        # nodes
+        mask_applied = mask_index == 0
+        xqb = x @ q_weights + q_bias
+        xqb_4d = xqb.reshape(shape_4d)
+        xkb = x @ k_weights + k_bias
+        xkb_4d = xkb.reshape(shape_4d)
+        xvb = x @ v_weights + v_bias
+        xvb_4d = xvb.reshape(shape_4d)
+        rot_xqb = np.transpose(xqb_4d, axes=(0, 2, 1, 3))
+        rot_xkb = np.transpose(xkb_4d, axes=(0, 2, 1, 3))
+        matmul = 0.125 * rot_xqb @ np.transpose(rot_xkb, [0, 1, 3, 2])
+        transpose_3 = np.transpose(xvb_4d, axes=(0, 2, 1, 3))
+        add_322 = matmul + attention_bias
+        masked_fill_2 = np.where(mask_applied, -np.inf, add_322)
+        softmax = scipy_special.softmax(masked_fill_2, axis=-1)
+        masked_fill_3 = np.where(mask_applied, 0, softmax)
+        matmul_1 = masked_fill_3 @ transpose_3
+        transpose_5 = np.transpose(matmul_1, axes=(0, 2, 1, 3))
+        view_3 = transpose_5.reshape(x.shape)
+        return (view_3,)

onnx_diagnostic/reference/ops/op_average_pool_grad.py

@@ -0,0 +1,63 @@
+import numpy as np
+from onnx.reference.op_run import OpRun
+
+
+class AveragePoolGrad(OpRun):
+    def _run(
+        self,
+        out,
+        auto_pad=None,
+        ceil_mode=None,
+        count_include_pad=None,
+        kernel_shape=None,
+        pads=None,
+        strides=None,
+    ):
+        assert auto_pad is not None, "auto_pad is None"
+        assert ceil_mode is not None, "ceil_mode is None"
+        assert count_include_pad is not None, "count_include_pad is None"
+        assert kernel_shape is not None, "kernel_shape is None"
+        assert pads is not None, "pads is None"
+        assert strides is not None, "strides is None"
+
+        assert auto_pad == "NOTSET", f"Not implemented for autopad={auto_pad!r}"
+        assert ceil_mode == 0, f"Not implemented for ceil_mode={ceil_mode!r}"
+        assert (
+            count_include_pad == 1
+        ), f"Not implemented for count_include_pad={count_include_pad!r}"
+
+        grad_shape = list(out.shape[:2])
+        for i in range(len(kernel_shape)):
+            d = (
+                out.shape[i + 2] * strides[i]
+                + kernel_shape[i]
+                - 1
+                + sum(pads[i * 2 : i * 2 + 2])
+            )
+            grad_shape.append(d)
+
+        grad = np.zeros(tuple(grad_shape), dtype=out.dtype)
+        scale = (1.0 / np.prod(kernel_shape)).astype(out.dtype)
+        if len(grad_shape) == 4:
+            # 2D
+            for batch in range(grad.shape[0]):
+                for channel in range(grad.shape[1]):
+                    for i in range(out.shape[2]):
+                        t = max(i * strides[0] - pads[0], 0)
+                        b = min(i * strides[0] - pads[0] + kernel_shape[0], grad.shape[2])
+                        for j in range(out.shape[3]):
+                            le = max(j * strides[1] - pads[2], 0)
+                            ri = min(
+                                j * strides[1] - pads[2] + kernel_shape[1],
+                                grad.shape[3],
+                            )
+
+                            grad[batch, channel, t:b, le:ri] += (
+                                out[batch, channel, i, j] * scale
+                            )
+        else:
+            raise NotImplementedError(
+                f"AveragePoolGrad is not implemented for shape={out.shape}."
+            )
+
+        return (grad.astype(out.dtype),)

onnx_diagnostic/reference/ops/op_bias_softmax.py

@@ -0,0 +1,16 @@
+import numpy as np
+from onnx.reference.op_run import OpRun
+
+
+class BiasSoftmax(OpRun):
+    op_domain = "com.microsoft"
+
+    def _run(self, x, y, axis=None, is_inner_broadcast=None):  # type: ignore
+        assert (
+            is_inner_broadcast == 0
+        ), f"Not implemented for is_inner_broadcast={is_inner_broadcast}"
+        z = x + y
+        tmp = z - z.max(axis=axis, keepdims=1)  # type: ignore
+        w = np.exp(tmp)
+        w /= w.sum(axis=axis, keepdims=1)  # type: ignore
+        return (w.astype(x.dtype),)

onnx_diagnostic/reference/ops/op_cast_like.py

@@ -0,0 +1,46 @@
+from onnx.onnx_pb import TensorProto
+from onnx.reference.op_run import OpRun
+
+try:
+    from onnx.reference.ops.op_cast import (
+        bfloat16,
+        cast_to,
+        float8e4m3fn,
+        float8e4m3fnuz,
+        float8e5m2,
+        float8e5m2fnuz,
+    )
+except ImportError:
+    bfloat16 = None
+    from onnx.reference.ops.op_cast import cast_to
+from ...helpers.onnx_helper import np_dtype_to_tensor_dtype
+
+
+def _cast_like(x, y, saturate):
+    if bfloat16 is not None:
+        if y.dtype == bfloat16 and y.dtype.descr[0][0] == "bfloat16":
+            # np.uint16 == np.uint16 is True as well as np.uint16 == bfloat16
+            to = TensorProto.BFLOAT16
+        elif y.dtype == float8e4m3fn and y.dtype.descr[0][0] == "e4m3fn":
+            to = TensorProto.FLOAT8E4M3FN
+        elif y.dtype == float8e4m3fnuz and y.dtype.descr[0][0] == "e4m3fnuz":
+            to = TensorProto.FLOAT8E4M3FNUZ
+        elif y.dtype == float8e5m2 and y.dtype.descr[0][0] == "e5m2":
+            to = TensorProto.FLOAT8E5M2
+        elif y.dtype == float8e5m2fnuz and y.dtype.descr[0][0] == "e5m2fnuz":
+            to = TensorProto.FLOAT8E5M2FNUZ
+        else:
+            to = np_dtype_to_tensor_dtype(y.dtype)  # type: ignore
+    else:
+        to = np_dtype_to_tensor_dtype(y.dtype)  # type: ignore
+    return (cast_to(x, to, saturate),)
+
+
+class CastLike_15(OpRun):
+    def _run(self, x, y):  # type: ignore
+        return _cast_like(x, y, True)
+
+
+class CastLike_19(OpRun):
+    def _run(self, x, y, saturate=None):  # type: ignore
+        return _cast_like(x, y, saturate)

onnx_diagnostic/reference/ops/op_complex.py

@@ -0,0 +1,26 @@
+import numpy as np
+from onnx.reference.op_run import OpRun
+
+
+class ToComplex(OpRun):
+    op_domain = "ai.onnx.complex"
+
+    def _run(self, x):
+        assert x.shape[-1] in (
+            1,
+            2,
+        ), f"Unexpected shape {x.shape}, it should a tensor (..., 2)"
+        if x.shape[-1] == 1:
+            return (x[..., 0] + 0j,)
+        return (x[..., 0] + 1j * x[..., 1],)
+
+
+class ComplexModule(OpRun):
+    op_domain = "ai.onnx.complex"
+
+    def _run(self, x):
+        assert x.dtype in (
+            np.complex64,
+            np.complex128,
+        ), f"Unexpected type {x.dtype}, it should a complex tensor"
+        return (np.abs(x),)

onnx_diagnostic/reference/ops/op_concat.py

@@ -0,0 +1,15 @@
+import numpy as np
+
+from onnx.reference.op_run import OpRun
+
+
+class Concat(OpRun):
+    def _preprocess(self, a: np.ndarray, axis: int) -> np.ndarray:
+        if axis >= len(a.shape):  # type: ignore
+            new_shape = a.shape + (1,) * (axis + 1 - len(a.shape))  # type: ignore
+            return a.reshape(new_shape)
+        return a
+
+    def _run(self, *args, axis=None):  # type: ignore
+        targs = tuple(self._preprocess(a, axis) for a in args)
+        return (np.concatenate(targs, axis),)  # type: ignore

onnx_diagnostic/reference/ops/op_constant_of_shape.py

@@ -0,0 +1,67 @@
+import numpy as np
+from onnx.reference.op_run import OpRun
+
+try:
+    import ml_dtypes
+except ImportError:
+    ml_dtypes = None  # type: ignore
+
+
+class ConstantOfShape(OpRun):
+    @staticmethod
+    def _process(value):
+        if (
+            value is not None
+            and ml_dtypes is not None
+            and value.dtype == (np.uint16, [("bfloat16", "<u2")])
+        ):
+            value = value.view(ml_dtypes.bfloat16)
+        cst = value[0] if isinstance(value, np.ndarray) and value.size > 0 else value
+        if isinstance(value, np.ndarray):
+            if not value.shape:
+                cst = value
+            elif value.size > 0:
+                cst = value.ravel()[0]
+            else:
+                raise ValueError(f"Unexpected fill_value={value!r}")
+        if isinstance(cst, bool):
+            cst = np.bool_(cst)
+        elif isinstance(cst, int):
+            cst = np.int64(cst)
+        elif isinstance(cst, float):
+            cst = np.float64(cst)
+        elif cst is None:
+            cst = np.float32(0)
+        if ml_dtypes is not None and isinstance(cst, ml_dtypes.bfloat16):
+            return cst
+        if not isinstance(
+            cst,
+            (
+                np.float16,
+                np.float32,
+                np.float64,
+                np.int64,
+                np.int32,
+                np.int16,
+                np.int8,
+                np.uint64,
+                np.uint32,
+                np.uint16,
+                np.uint8,
+                np.bool_,
+            ),
+        ):
+            raise TypeError(f"value must be a real not {type(cst)}")
+        return cst
+
+    def _run(self, data, value=None):
+        cst = self._process(value)
+        try:
+            res = np.full(tuple(data), cst)
+        except TypeError as e:
+            raise RuntimeError(
+                f"Unable to create a constant of shape "
+                f"{data!r} with value {cst!r} "
+                f"(raw value={value!r})."
+            ) from e
+        return (res,)

onnx_diagnostic/reference/ops/op_fused_matmul.py

@@ -0,0 +1,31 @@
+import numpy as np
+from onnx.reference.op_run import OpRun
+
+
+class FusedMatMul(OpRun):
+    op_domain = "com.microsoft"
+
+    def _run(
+        self,
+        A,
+        B,
+        alpha: float = 1,
+        transA: int = 0,
+        transB: int = 0,
+        transBatchA: int = 0,
+        transBatchB: int = 0,
+    ):
+        assert transBatchA == 0, f"Not implemented for transBatchA==1 and {A.shape}x{B.shape}"
+        assert transBatchB == 0, f"Not implemented for transBatchB==1 and {A.shape}x{B.shape}"
+        if transA:
+            perm = list(range(len(A.shape)))
+            dim = len(perm)
+            perm[dim - 2], perm[dim - 1] = perm[dim - 1], perm[dim - 2]
+            A = np.transpose(A, perm)
+        if transB:
+            perm = list(range(len(B.shape)))
+            dim = len(perm)
+            perm[dim - 2], perm[dim - 1] = perm[dim - 1], perm[dim - 2]
+            B = np.transpose(B, perm)
+        a = np.array(alpha, dtype=A.dtype)
+        return (np.matmul(A, B) * a,)

onnx_diagnostic/reference/ops/op_gather.py

@@ -0,0 +1,29 @@
+# Copyright (c) ONNX Project Contributors
+
+# SPDX-License-Identifier: Apache-2.0
+from __future__ import annotations
+
+import numpy as np
+
+from onnx.reference.op_run import OpRun
+
+
+class Gather(OpRun):
+    def _run(self, x, indices, axis=None):
+        if x.size == 0 or indices.size == 0:
+            if axis is None:
+                new_shape = indices.shape
+            else:
+                new_shape = (*x.shape[:axis], *indices.shape, *x.shape[axis + 1 :])
+            if 0 not in new_shape:
+                new_shape = (0, *new_shape[1:])
+            return (np.empty(new_shape, dtype=x.dtype),)
+        if not x.flags["C_CONTIGUOUS"]:
+            x = np.ascontiguousarray(x)
+        if not indices.flags["C_CONTIGUOUS"]:
+            indices = indices.ascontiguousarray()
+        try:
+            return (np.take(x, indices, axis=axis),)
+        except TypeError:
+            # distribution x86 requires int32.
+            return (np.take(x, indices.astype(int), axis=axis),)