emx-onnx-cgen 0.2.0__py3-none-any.whl → 0.3.1__py3-none-any.whl

emx_onnx_cgen/lowering/shape.py

@@ -2,32 +2,13 @@ from __future__ import annotations
 
 from shared.scalar_types import ScalarType
 
-from ..codegen.c_emitter import ShapeOp
+from ..ir.ops import ShapeOp
 from ..errors import ShapeInferenceError, UnsupportedOpError
 from ..ir.model import Graph, Node
+from .common import value_dtype, value_shape
 from .registry import register_lowering
 
 
-def _value_shape(graph: Graph, name: str, node: Node) -> tuple[int, ...]:
-    try:
-        return graph.find_value(name).type.shape
-    except KeyError as exc:
-        raise ShapeInferenceError(
-            f"Missing shape for value '{name}' in op {node.op_type}. "
-            "Hint: run ONNX shape inference or export with static shapes."
-        ) from exc
-
-
-def _value_dtype(graph: Graph, name: str, node: Node) -> ScalarType:
-    try:
-        return graph.find_value(name).type.dtype
-    except KeyError as exc:
-        raise ShapeInferenceError(
-            f"Missing dtype for value '{name}' in op {node.op_type}. "
-            "Hint: run ONNX shape inference or export with static shapes."
-        ) from exc
-
-
 def _normalize_slice_bounds(
     rank: int, *, start: int | None, end: int | None
 ) -> tuple[int, int]:
@@ -46,14 +27,14 @@ def _normalize_slice_bounds(
 def lower_shape(graph: Graph, node: Node) -> ShapeOp:
     if len(node.inputs) != 1 or len(node.outputs) != 1:
         raise UnsupportedOpError("Shape must have 1 input and 1 output")
-    input_shape = _value_shape(graph, node.inputs[0], node)
-    output_shape = _value_shape(graph, node.outputs[0], node)
+    input_shape = value_shape(graph, node.inputs[0], node)
+    output_shape = value_shape(graph, node.outputs[0], node)
     if len(output_shape) != 1:
         raise ShapeInferenceError("Shape output must be 1D")
     if output_shape[0] < 0:
         raise ShapeInferenceError("Shape output length must be non-negative")
-    input_dtype = _value_dtype(graph, node.inputs[0], node)
-    output_dtype = _value_dtype(graph, node.outputs[0], node)
+    input_dtype = value_dtype(graph, node.inputs[0], node)
+    output_dtype = value_dtype(graph, node.outputs[0], node)
     if output_dtype != ScalarType.I64:
         raise UnsupportedOpError("Shape output dtype must be int64")
     start = node.attrs.get("start")

emx_onnx_cgen/lowering/size.py

@@ -2,7 +2,7 @@ from __future__ import annotations
 
 from shared.scalar_types import ScalarType
 
-from ..codegen.c_emitter import SizeOp
+from ..ir.ops import SizeOp
 from ..errors import ShapeInferenceError, UnsupportedOpError
 from ..ir.model import Graph, Node
 from .common import shape_product, value_dtype, value_shape

emx_onnx_cgen/lowering/slice.py

@@ -6,7 +6,7 @@ import numpy as np
 
 from shared.scalar_types import ScalarType
 
-from ..codegen.c_emitter import SliceOp
+from ..ir.ops import SliceOp
 from ..errors import ShapeInferenceError, UnsupportedOpError
 from ..ir.model import Graph, Initializer, Node
 from ..lowering.common import value_dtype, value_shape

emx_onnx_cgen/lowering/softmax.py

@@ -1,9 +1,10 @@
 from __future__ import annotations
 
-from ..codegen.c_emitter import SoftmaxOp
+from ..ir.ops import SoftmaxOp
 from ..errors import UnsupportedOpError
 from ..ir.model import Graph, Node
 from .common import node_dtype as _node_dtype
+from .common import onnx_opset_version as _onnx_opset_version
 from .common import shape_product as _shape_product
 from .common import value_shape as _value_shape
 from .registry import register_lowering
@@ -23,8 +24,11 @@ def lower_softmax(graph: Graph, node: Node) -> SoftmaxOp:
     input_shape = _value_shape(graph, node.inputs[0], node)
     output_shape = _value_shape(graph, node.outputs[0], node)
     ensure_output_shape_matches_input(node, input_shape, output_shape)
+    opset_version = _onnx_opset_version(graph)
+    default_axis = 1 if opset_version is not None and opset_version < 13 else -1
+    axis_attr = node.attrs.get("axis", default_axis)
     axis = _normalize_axis(
-        int(node.attrs.get("axis", -1)),
+        int(axis_attr),
         input_shape,
         node,
     )

emx_onnx_cgen/lowering/softmax_cross_entropy_loss.py

@@ -2,7 +2,7 @@ from __future__ import annotations
 
 from shared.scalar_types import ScalarType
 
-from ..codegen.c_emitter import SoftmaxCrossEntropyLossOp
+from ..ir.ops import SoftmaxCrossEntropyLossOp
 from ..errors import ShapeInferenceError, UnsupportedOpError
 from ..ir.model import Graph, Node
 from .common import shape_product as _shape_product

emx_onnx_cgen/lowering/split.py

@@ -4,7 +4,7 @@ import numpy as np
 
 from shared.scalar_types import ScalarType
 
-from ..codegen.c_emitter import SplitOp
+from ..ir.ops import SplitOp
 from ..errors import ShapeInferenceError, UnsupportedOpError
 from ..ir.model import Graph, Initializer, Node
 from ..lowering.common import optional_name, value_dtype, value_shape

emx_onnx_cgen/lowering/squeeze.py

@@ -2,7 +2,7 @@ from __future__ import annotations
 
 from shared.scalar_types import ScalarType
 
-from ..codegen.c_emitter import ReshapeOp
+from ..ir.ops import ReshapeOp
 from ..errors import ShapeInferenceError, UnsupportedOpError
 from ..ir.model import Graph, Initializer, Node
 from .registry import register_lowering
@@ -95,11 +95,11 @@ def _validate_output_shape_for_unknown_axes(
     for dim in input_shape:
         if output_index < len(output_shape) and dim == output_shape[output_index]:
             output_index += 1
-            continue
-        if dim != 1:
-            raise ShapeInferenceError(
-                "Squeeze output shape must remove only dimensions of size 1"
-            )
+        else:
+            if dim != 1:
+                raise ShapeInferenceError(
+                    "Squeeze output shape must remove only dimensions of size 1"
+                )
     if output_index != len(output_shape):
         raise ShapeInferenceError(
             "Squeeze output shape must preserve input order while removing size-1 axes"

emx_onnx_cgen/lowering/tensor_scatter.py

@@ -0,0 +1,110 @@
+from __future__ import annotations
+
+from shared.scalar_types import ScalarType
+
+from ..ir.ops import TensorScatterOp
+from ..errors import ShapeInferenceError, UnsupportedOpError
+from ..ir.model import Graph, Node
+from ..validation import normalize_axis
+from .common import optional_name, value_dtype, value_shape
+from .registry import register_lowering
+
+_ALLOWED_MODES = {"linear", "circular"}
+
+
+@register_lowering("TensorScatter")
+def lower_tensor_scatter(graph: Graph, node: Node) -> TensorScatterOp:
+    if len(node.inputs) not in {2, 3} or len(node.outputs) != 1:
+        raise UnsupportedOpError(
+            "TensorScatter must have 2 or 3 inputs and 1 output"
+        )
+    past_cache_name = node.inputs[0]
+    update_name = node.inputs[1]
+    write_indices_name = optional_name(node.inputs, 2)
+    output_name = node.outputs[0]
+    past_cache_shape = value_shape(graph, past_cache_name, node)
+    update_shape = value_shape(graph, update_name, node)
+    output_shape = value_shape(graph, output_name, node)
+    if output_shape != past_cache_shape:
+        raise ShapeInferenceError(
+            "TensorScatter output shape must match past_cache shape, "
+            f"got {output_shape} vs {past_cache_shape}"
+        )
+    if len(update_shape) != len(past_cache_shape):
+        raise ShapeInferenceError(
+            "TensorScatter update shape rank must match past_cache rank, "
+            f"got {len(update_shape)} vs {len(past_cache_shape)}"
+        )
+    axis = normalize_axis(int(node.attrs.get("axis", -2)), past_cache_shape, node)
+    if axis == 0:
+        raise UnsupportedOpError(
+            "TensorScatter axis cannot be 0 (batch dimension)"
+        )
+    for dim_index, (past_dim, update_dim) in enumerate(
+        zip(past_cache_shape, update_shape)
+    ):
+        if dim_index == axis:
+            if update_dim > past_dim:
+                raise ShapeInferenceError(
+                    "TensorScatter update sequence length must be <= "
+                    "past_cache sequence length, "
+                    f"got {update_dim} vs {past_dim}"
+                )
+        elif update_dim != past_dim:
+            raise ShapeInferenceError(
+                "TensorScatter update shape must match past_cache shape "
+                f"outside axis {axis}, got {update_shape} vs {past_cache_shape}"
+            )
+    mode = node.attrs.get("mode", "linear")
+    if isinstance(mode, bytes):
+        mode = mode.decode("utf-8")
+    if mode not in _ALLOWED_MODES:
+        raise UnsupportedOpError(
+            "TensorScatter mode must be one of "
+            f"{sorted(_ALLOWED_MODES)}, got {mode}"
+        )
+    dtype = value_dtype(graph, past_cache_name, node)
+    update_dtype = value_dtype(graph, update_name, node)
+    output_dtype = value_dtype(graph, output_name, node)
+    if update_dtype != dtype or output_dtype != dtype:
+        raise UnsupportedOpError(
+            "TensorScatter expects past_cache, update, and output "
+            "to share the same dtype, "
+            f"got {dtype.onnx_name}, {update_dtype.onnx_name}, "
+            f"{output_dtype.onnx_name}"
+        )
+    write_indices_shape = None
+    write_indices_dtype = None
+    if write_indices_name is not None:
+        write_indices_shape = value_shape(graph, write_indices_name, node)
+        if len(write_indices_shape) != 1:
+            raise ShapeInferenceError(
+                "TensorScatter write_indices must be a 1D tensor"
+            )
+        if write_indices_shape[0] != past_cache_shape[0]:
+            raise ShapeInferenceError(
+                "TensorScatter write_indices length must match batch size, "
+                f"got {write_indices_shape[0]} vs {past_cache_shape[0]}"
+            )
+        write_indices_dtype = value_dtype(
+            graph, write_indices_name, node
+        )
+        if write_indices_dtype not in {ScalarType.I64, ScalarType.I32}:
+            raise UnsupportedOpError(
+                "TensorScatter write_indices must be int32 or int64, "
+                f"got {write_indices_dtype.onnx_name}"
+            )
+    return TensorScatterOp(
+        past_cache=past_cache_name,
+        update=update_name,
+        write_indices=write_indices_name,
+        output=output_name,
+        past_cache_shape=past_cache_shape,
+        update_shape=update_shape,
+        output_shape=output_shape,
+        write_indices_shape=write_indices_shape,
+        axis=axis,
+        mode=mode,
+        dtype=dtype,
+        write_indices_dtype=write_indices_dtype,
+    )

emx_onnx_cgen/lowering/tile.py

@@ -4,7 +4,7 @@ import numpy as np
 
 from shared.scalar_types import ScalarType
 
-from ..codegen.c_emitter import TileOp
+from ..ir.ops import TileOp
 from ..errors import ShapeInferenceError, UnsupportedOpError
 from ..ir.model import Graph, Initializer, Node
 from ..lowering.common import value_dtype, value_shape

emx_onnx_cgen/lowering/topk.py

@@ -0,0 +1,134 @@
+from __future__ import annotations
+
+import numpy as np
+
+from shared.scalar_types import ScalarType
+
+from ..ir.ops import TopKOp
+from ..errors import ShapeInferenceError, UnsupportedOpError
+from ..ir.model import Graph, Initializer, Node
+from ..lowering.common import shape_product, value_dtype, value_shape
+from ..validation import normalize_axis
+from .registry import register_lowering
+
+
+def _find_initializer(graph: Graph, name: str) -> Initializer | None:
+    for initializer in graph.initializers:
+        if initializer.name == name:
+            return initializer
+    return None
+
+
+def _read_k(graph: Graph, name: str, node: Node) -> int | None:
+    initializer = _find_initializer(graph, name)
+    if initializer is None:
+        return None
+    if initializer.type.dtype not in {ScalarType.I64, ScalarType.I32}:
+        raise UnsupportedOpError(
+            f"{node.op_type} k input must be int64 or int32"
+        )
+    data = np.array(initializer.data, dtype=np.int64).reshape(-1)
+    if data.size != 1:
+        raise ShapeInferenceError(
+            f"{node.op_type} k input must contain a single value"
+        )
+    k = int(data[0])
+    if k <= 0:
+        raise ShapeInferenceError(
+            f"{node.op_type} k must be a positive value, got {k}"
+        )
+    return k
+
+
+def _topk_dtype_supported(dtype: ScalarType) -> bool:
+    return not dtype.is_bool
+
+
+def lower_topk(graph: Graph, node: Node) -> TopKOp:
+    if node.op_type != "TopK":
+        raise UnsupportedOpError(f"Unsupported op {node.op_type}")
+    if len(node.inputs) != 2 or len(node.outputs) != 2:
+        raise UnsupportedOpError(
+            f"{node.op_type} must have 2 inputs and 2 outputs"
+        )
+    input_name = node.inputs[0]
+    k_name = node.inputs[1]
+    output_values = node.outputs[0]
+    output_indices = node.outputs[1]
+    input_shape = value_shape(graph, input_name, node)
+    shape_product(input_shape)
+    axis = int(node.attrs.get("axis", -1))
+    axis = normalize_axis(axis, input_shape, node)
+    k = _read_k(graph, k_name, node)
+    axis_dim = input_shape[axis]
+    values_shape = value_shape(graph, output_values, node)
+    indices_shape = value_shape(graph, output_indices, node)
+    if values_shape != indices_shape:
+        raise ShapeInferenceError(
+            f"{node.op_type} values and indices output shapes must match, "
+            f"got {values_shape} and {indices_shape}"
+        )
+    if k is None:
+        k_shape = value_shape(graph, k_name, node)
+        if len(k_shape) != 1 or k_shape[0] != 1:
+            raise ShapeInferenceError(
+                f"{node.op_type} k input must be a 1-element tensor"
+            )
+        if axis >= len(values_shape):
+            raise ShapeInferenceError(
+                f"{node.op_type} axis {axis} exceeds output rank {len(values_shape)}"
+            )
+        k = values_shape[axis]
+        if k <= 0:
+            raise ShapeInferenceError(
+                f"{node.op_type} k must be a positive value, got {k}"
+            )
+    if k > axis_dim:
+        raise ShapeInferenceError(
+            f"{node.op_type} k {k} exceeds axis dimension {axis_dim}"
+        )
+    output_shape_expected = list(input_shape)
+    output_shape_expected[axis] = k
+    output_shape = tuple(output_shape_expected)
+    if values_shape != output_shape:
+        raise ShapeInferenceError(
+            f"{node.op_type} values output shape must be {output_shape}, got {values_shape}"
+        )
+    if indices_shape != output_shape:
+        raise ShapeInferenceError(
+            f"{node.op_type} indices output shape must be {output_shape}, got {indices_shape}"
+        )
+    input_dtype = value_dtype(graph, input_name, node)
+    if not _topk_dtype_supported(input_dtype):
+        raise UnsupportedOpError(
+            f"{node.op_type} does not support dtype {input_dtype.onnx_name}"
+        )
+    values_dtype = value_dtype(graph, output_values, node)
+    if values_dtype != input_dtype:
+        raise UnsupportedOpError(
+            f"{node.op_type} values output dtype must be {input_dtype.onnx_name}"
+        )
+    indices_dtype = value_dtype(graph, output_indices, node)
+    if indices_dtype != ScalarType.I64:
+        raise UnsupportedOpError(
+            f"{node.op_type} indices output dtype must be int64"
+        )
+    largest = bool(int(node.attrs.get("largest", 1)))
+    sorted_output = bool(int(node.attrs.get("sorted", 1)))
+    return TopKOp(
+        input0=input_name,
+        output_values=output_values,
+        output_indices=output_indices,
+        input_shape=input_shape,
+        output_shape=output_shape,
+        axis=axis,
+        k=k,
+        largest=largest,
+        sorted=sorted_output,
+        input_dtype=input_dtype,
+        output_values_dtype=values_dtype,
+        output_indices_dtype=indices_dtype,
+    )
+
+
+register_lowering("TopK")(lower_topk)

emx_onnx_cgen/lowering/transpose.py

@@ -1,6 +1,6 @@
 from __future__ import annotations
 
-from ..codegen.c_emitter import TransposeOp
+from ..ir.ops import TransposeOp
 from ..errors import ShapeInferenceError, UnsupportedOpError
 from ..ir.model import Graph, Node
 from .common import node_dtype as _node_dtype

emx_onnx_cgen/lowering/trilu.py

@@ -0,0 +1,89 @@
+from __future__ import annotations
+
+import numpy as np
+
+from shared.scalar_types import ScalarType
+
+from ..ir.ops import TriluOp
+from ..errors import ShapeInferenceError, UnsupportedOpError
+from ..ir.model import Graph, Initializer, Node
+from ..lowering.common import optional_name, value_dtype, value_shape
+from .registry import register_lowering
+
+
+def _find_initializer(graph: Graph, name: str) -> Initializer | None:
+    for initializer in graph.initializers:
+        if initializer.name == name:
+            return initializer
+    return None
+
+
+def _is_scalar_shape(shape: tuple[int, ...]) -> bool:
+    return shape == () or shape == (1,)
+
+
+def _read_k_initializer(initializer: Initializer, node: Node) -> int:
+    if initializer.type.dtype != ScalarType.I64:
+        raise UnsupportedOpError(
+            f"{node.op_type} k input must be int64"
+        )
+    data = np.array(initializer.data, dtype=np.int64).reshape(-1)
+    if data.size != 1:
+        raise UnsupportedOpError(f"{node.op_type} k input must be scalar")
+    return int(data[0])
+
+
+@register_lowering("Trilu")
+def lower_trilu(graph: Graph, node: Node) -> TriluOp:
+    if len(node.inputs) not in {1, 2} or len(node.outputs) != 1:
+        raise UnsupportedOpError("Trilu must have 1 or 2 inputs and 1 output")
+    input_name = node.inputs[0]
+    if not input_name:
+        raise UnsupportedOpError("Trilu input must be provided")
+    input_shape = value_shape(graph, input_name, node)
+    output_shape = value_shape(graph, node.outputs[0], node)
+    if input_shape != output_shape:
+        raise ShapeInferenceError("Trilu input and output shapes must match")
+    if len(output_shape) < 2:
+        raise UnsupportedOpError("Trilu expects input rank >= 2")
+    input_dtype = value_dtype(graph, input_name, node)
+    output_dtype = value_dtype(graph, node.outputs[0], node)
+    if input_dtype != output_dtype:
+        raise UnsupportedOpError(
+            "Trilu expects matching input/output dtypes, "
+            f"got {input_dtype.onnx_name} and {output_dtype.onnx_name}"
+        )
+    upper_attr = node.attrs.get("upper", 1)
+    upper = bool(int(upper_attr))
+    k_input = optional_name(node.inputs, 1)
+    k_value = 0
+    k_input_name = None
+    k_input_shape = None
+    k_input_dtype = None
+    if k_input:
+        k_initializer = _find_initializer(graph, k_input)
+        if k_initializer is not None:
+            k_value = _read_k_initializer(k_initializer, node)
+        else:
+            k_shape = value_shape(graph, k_input, node)
+            if not _is_scalar_shape(k_shape):
+                raise UnsupportedOpError("Trilu k input must be scalar")
+            k_dtype = value_dtype(graph, k_input, node)
+            if k_dtype != ScalarType.I64:
+                raise UnsupportedOpError("Trilu k input must be int64")
+            k_input_name = k_input
+            k_input_shape = k_shape
+            k_input_dtype = k_dtype
+    return TriluOp(
+        input0=input_name,
+        output=node.outputs[0],
+        input_shape=input_shape,
+        output_shape=output_shape,
+        upper=upper,
+        k_value=k_value,
+        k_input=k_input_name,
+        k_input_shape=k_input_shape,
+        k_input_dtype=k_input_dtype,
+        dtype=output_dtype,
+        input_dtype=input_dtype,
+    )

emx_onnx_cgen/lowering/unsqueeze.py

@@ -2,7 +2,7 @@ from __future__ import annotations
 
 from shared.scalar_types import ScalarType
 
-from ..codegen.c_emitter import ReshapeOp
+from ..ir.ops import ReshapeOp
 from ..errors import ShapeInferenceError, UnsupportedOpError
 from ..ir.model import Graph, Initializer, Node
 from .registry import register_lowering
@@ -131,11 +131,11 @@ def lower_unsqueeze(graph: Graph, node: Node) -> ReshapeOp:
         for dim in output_shape:
             if input_index < len(input_shape) and dim == input_shape[input_index]:
                 input_index += 1
-                continue
-            if dim != 1:
-                raise ShapeInferenceError(
-                    "Unsqueeze output shape must insert ones only"
-                )
+            else:
+                if dim != 1:
+                    raise ShapeInferenceError(
+                        "Unsqueeze output shape must insert ones only"
+                    )
         if input_index != len(input_shape):
             raise ShapeInferenceError(
                 "Unsqueeze output shape must contain input shape in order"

emx_onnx_cgen/lowering/variadic.py

@@ -3,7 +3,7 @@ from __future__ import annotations
 from shared.scalar_functions import ScalarFunction
 from shared.scalar_types import ScalarType
 
-from ..codegen.c_emitter import MultiInputBinaryOp
+from ..ir.ops import MultiInputBinaryOp
 from ..errors import UnsupportedOpError
 from ..ir.model import Graph, Node
 from ..lowering.common import node_dtype, value_dtype, value_shape

emx_onnx_cgen/lowering/where.py

@@ -2,7 +2,7 @@ from __future__ import annotations
 
 from shared.scalar_types import ScalarType
 
-from ..codegen.c_emitter import WhereOp
+from ..ir.ops import WhereOp
 from ..errors import ShapeInferenceError, UnsupportedOpError
 from ..ir.model import Graph, Node
 from .common import value_dtype as _value_dtype

emx_onnx_cgen/onnx_import.py

@@ -212,6 +212,9 @@ def import_onnx(model: onnx.ModelProto) -> Graph:
     dim_param_by_name = _collect_dim_params(
         tuple(model.graph.input) + tuple(model.graph.output)
     )
+    opset_imports = tuple(
+        (opset.domain, opset.version) for opset in model.opset_import
+    )
     try:
         model = shape_inference.infer_shapes(model, data_prop=True)
     except Exception as exc:  # pragma: no cover - onnx inference errors
@@ -258,4 +261,5 @@ def import_onnx(model: onnx.ModelProto) -> Graph:
         nodes=nodes,
         initializers=initializers,
         values=values,
+        opset_imports=opset_imports,
     )

emx_onnx_cgen/onnxruntime_utils.py

@@ -0,0 +1,11 @@
+from __future__ import annotations
+
+from typing import Any
+
+
+def make_deterministic_session_options(ort: Any) -> Any:
+    options = ort.SessionOptions()
+    options.intra_op_num_threads = 1
+    options.inter_op_num_threads = 1
+    options.execution_mode = ort.ExecutionMode.ORT_SEQUENTIAL
+    return options

emx_onnx_cgen/ops.py

@@ -87,6 +87,7 @@ UNARY_OP_TYPES = {
     "Identity",
     "LeakyRelu",
     "Log",
+    "Mish",
     "Neg",
     "Not",
     "Reciprocal",
@@ -177,6 +178,7 @@ UNARY_SYMBOLS_DOUBLE = {
     ScalarFunction.LEAKY_RELU: "leaky_relu",
     ScalarFunction.POSITIVE: "identity",
     ScalarFunction.LOG: "log",
+    ScalarFunction.MISH: "mish",
     ScalarFunction.NEG: "neg",
     ScalarFunction.RECIPROCAL: "reciprocal",
     ScalarFunction.RELU: "relu",
@@ -215,6 +217,7 @@ UNARY_SYMBOLS_FLOAT = {
     ScalarFunction.LEAKY_RELU: "leaky_relu",
     ScalarFunction.POSITIVE: "identity",
     ScalarFunction.LOG: "logf",
+    ScalarFunction.MISH: "mish",
     ScalarFunction.NEG: "neg",
     ScalarFunction.RECIPROCAL: "reciprocal",
     ScalarFunction.RELU: "relu",
@@ -457,6 +460,7 @@ UNARY_APPLY_FUNCS = {
     "thresholded_relu": lambda value: np.where(
         value > 1.0, value, 0.0
     ),
+    "mish": lambda value: value * np.tanh(np.log1p(np.exp(value))),
     "atanhf": np.arctanh,
     "atanh": np.arctanh,
 }