torchax 0.0.4__py3-none-any.whl

torchax/ops/mappings.py

@@ -0,0 +1,97 @@
+from jax import dlpack as jaxdl
+import jax.numpy as jnp
+import numpy
+import torch
+import torch.func
+import torch.utils.dlpack as torchdl
+import torch.utils._mode_utils as mode_utils
+
+
+def t2j(t):
+  is_bool = False
+  if t.dtype == torch.bool:
+    is_bool = True
+    t = t.to(torch.int8)
+
+  t = t.to_dense()
+
+  if not t.is_contiguous():
+    t = t.contiguous()
+
+  try:
+    res = jaxdl.from_dlpack(t)
+  except Exception:
+    # https://github.com/google/jax/issues/7657
+    # https://github.com/google/jax/issues/17784
+    if t.dtype == torch.bfloat16:
+      nparray = (t.cpu().detach().to(torch.float32).numpy()
+                )  # numpy don't support bfloat16
+    else:
+      nparray = t.cpu().detach().numpy()
+    res = jnp.asarray(nparray)
+    if t.dtype == torch.bfloat16:
+      res = res.astype(jnp.bfloat16)
+
+  if is_bool:
+    res = res.astype(jnp.bool_)
+  return res
+
+
+def j2t(x):
+  with mode_utils.no_dispatch(), torch._C.DisableTorchFunction():
+    try:
+      dl = jaxdl.to_dlpack(x)
+      res = torchdl.from_dlpack(dl)
+    except Exception:
+      res = torch.from_numpy(numpy.asarray(x))
+    if x.dtype == jnp.bool_:
+      res = res.to(torch.bool)
+    return res
+
+TORCH_DTYPE_TO_JAX = {
+    # NO_MAPPING        : jnp.float0.dtype (signless scalar int),
+    torch.bool          : jnp.bool_.dtype,
+    # NO_MAPPING        : jnp.int4.dtype,
+    torch.int8          : jnp.int8.dtype,
+    torch.int16         : jnp.int16.dtype,
+    torch.int32         : jnp.int32.dtype,
+    torch.int64         : jnp.int64.dtype,
+    torch.long          : jnp.int64.dtype,
+    # NO_MAPPING        : jnp.uint4
+    torch.uint8         : jnp.uint8.dtype,
+    torch.uint16        : jnp.uint16.dtype,
+    torch.uint32        : jnp.uint32.dtype,
+    torch.uint64        : jnp.uint64.dtype,
+    # NO_MAPPING        : jnp.float8_e4m3b11fnuz.dtype,
+    torch.float8_e4m3fn : jnp.float8_e4m3fn.dtype,
+    # NO_MAPPING        : jnp.float8_e4m3fnuz.dtype,
+    torch.float8_e5m2   : jnp.float8_e5m2.dtype,
+    # NO_MAPPING        : jnp.float8_e5m2fnuz.dtype,
+    torch.bfloat16      : jnp.bfloat16.dtype,
+    torch.half          : jnp.float16.dtype,
+    torch.float16       : jnp.float16.dtype,
+    torch.float32       : jnp.float32.dtype,
+    torch.float64       : jnp.float64.dtype,
+    torch.double        : jnp.double.dtype,
+    torch.complex64     : jnp.complex64.dtype,
+    torch.complex128    : jnp.complex128.dtype,
+    None                : None,
+}
+
+JAX_DTYPE_TO_TORCH = {
+  value: key for key, value in TORCH_DTYPE_TO_JAX.items()
+}
+# Add imprecise mappings for some JAX dtypes which don't have torch analogues
+JAX_DTYPE_TO_TORCH[jnp.dtype('int4')] = torch.int8
+JAX_DTYPE_TO_TORCH[jnp.dtype('uint4')] = torch.uint8
+
+def t2j_dtype(dtype):
+  if dtype not in TORCH_DTYPE_TO_JAX:
+    raise RuntimeError(f'Attempting to convert unknown type: {dtype} to jax type,')
+  return TORCH_DTYPE_TO_JAX[dtype]
+
+
+def j2t_dtype(dtype):
+  if dtype not in JAX_DTYPE_TO_TORCH:
+    raise RuntimeError(f'Attempting to convert unknown type: {dtype} to torch type,')
+  return JAX_DTYPE_TO_TORCH[dtype]

torchax/ops/op_base.py

@@ -0,0 +1,104 @@
+import functools
+import jax
+import jax.numpy as jnp
+import numpy as np
+import torch
+from torchax.ops import mappings
+from torchax import types
+import sys
+
+from typing import Callable, Optional, ParamSpec, Concatenate
+
+
+class InplaceOp:
+
+    def __init__(self, functional_op, replace=False, position_to_mutate=0):
+        self.functional = functional_op
+        self.replace = replace
+        self.position_to_mutate = position_to_mutate
+
+    def __call__(self, *args, **kwargs):
+        to_mutate = args[0]
+        if self.replace:
+          to_mutate._elem = self.functional(*args, **kwargs)._elem
+        else:
+          to_mutate.copy_(self.functional(*args, **kwargs))
+        return to_mutate
+
+
+class OutVariant:
+
+    def __call__(self, *args, **kwargs):
+        to_mutate = kwargs['out']
+        del kwargs['out']
+        to_mutate._elem = self.functional(*args, **kwargs)._elem
+        return to_mutate
+
+
+
+P = ParamSpec('P')
+def convert_dtype(use_default_dtype: bool = True):
+  """Converts `dtype` kwarg of function from torch to JAX.
+
+  Args:
+    use_default_dtype: Whether to use torch default dtype if none is provided.
+
+  Returns:
+    A decorator that wraps a JAX implementation of a torch function.
+  """
+
+  def decorator(func: types.TorchCallable):
+    @functools.wraps(func)
+    def wrapper(*args: P.args,
+                dtype: Optional[torch.dtype] = None,
+                **kwargs: P.kwargs):
+      if not dtype and use_default_dtype:
+        dtype = torch.get_default_dtype()
+      if isinstance(dtype, torch.dtype):
+        jax_dtype = mappings.t2j_dtype(dtype)
+      else:
+        jax_dtype = dtype
+
+      return func(*args, dtype=jax_dtype, **kwargs)
+
+    return wrapper
+
+  return decorator
+
+
+def maybe_convert_constant_dtype(val: Optional[types.JaxValue], dtype: Optional[jnp.dtype]):
+  """Optionally converts scalar constant's dtype using `numpy`
+
+  Use in cases where you require a constant and can't handle a traced array.
+  """
+  if val and dtype:
+    if isinstance(val, jax.Array):
+      return maybe_convert_constant_dtype(val.item(), dtype)
+
+    return np.array(val, dtype)
+
+  return val
+
+
+def promote_int_input(f: Callable[Concatenate[jax.Array, P], types.JaxValue]):
+   """If the first argument is an int array, promote it to float32."""
+   @functools.wraps(f)
+   def wrapper(x: jax.Array, *args: P.args, **kwargs: P.kwargs):
+      if x.dtype in [jnp.int8, jnp.int16, jnp.int32, jnp.int64]:
+        x = x.astype(mappings.t2j_dtype(torch.get_default_dtype()))
+
+      return f(x, *args, **kwargs)
+
+   return wrapper
+
+
+def foreach_loop(
+  seq: jax.Array, fn: Callable[[jax.Array, jax.Array], jax.Array], init_val=0.0
+):
+  """Run `fn` for each element of 1D array `seq`.
+
+  Similar to `functools.reduce`, but implemented with `jax.lax.fori_loop`."""
+  assert len(seq.shape) == 1
+  return jax.lax.fori_loop(
+    0, len(seq), lambda i, carry: fn(carry, seq[i]), init_val
+  )

torchax/ops/ops_registry.py

@@ -0,0 +1,50 @@
+import dataclasses
+import logging
+from torchax.types import JaxCallable, TorchCallable
+
+from typing import Union, Dict
+
+
+@dataclasses.dataclass
+class Operator:
+    torch_op: TorchCallable
+    func: Union[TorchCallable, JaxCallable]
+    is_jax_function: bool
+    is_user_defined: bool
+    needs_env: bool
+
+
+all_aten_ops: Dict[TorchCallable, Operator] = {}
+all_torch_functions: Dict[TorchCallable, Operator] = {}
+
+
+def register_torch_dispatch_op(
+    aten_op, impl_callable,
+    is_jax_function=True,
+    is_user_defined=False,
+    needs_env=False,
+):
+    op = Operator(
+        aten_op, impl_callable,
+        is_jax_function=is_jax_function,
+        is_user_defined=is_user_defined,
+        needs_env=needs_env)
+    if aten_op in all_aten_ops:
+        logging.warning(f'Duplicate op registration for {aten_op}')
+    all_aten_ops[aten_op] = op
+    return impl_callable
+
+
+def register_torch_function_op(
+    torch_func, impl_callable,
+    is_jax_function=True,
+    is_user_defined=False,
+    needs_env=False,
+):
+    op = Operator(
+        torch_func, impl_callable,
+        is_jax_function=is_jax_function,
+        is_user_defined=is_user_defined,
+        needs_env=needs_env)
+    all_torch_functions[torch_func] = op
+    return impl_callable