# Copyright 2022 The JAX Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Module for state primitives."""
from __future__ import annotations

from functools import partial
import json
import types
from typing import Any

import numpy as np

from jax._src import ad_util
from jax._src import core
from jax._src import dispatch
from jax._src import dtypes
from jax._src import pretty_printer as pp
from jax._src import source_info_util
from jax._src import traceback_util
from jax._src import tree_util
from jax._src.interpreters import ad
from jax._src.interpreters import batching
from jax._src.interpreters import mlir
from jax._src.interpreters import partial_eval as pe
from jax._src.interpreters import remat
from jax._src.lax import lax
from jax._src.state import indexing
from jax._src.state.types import (
    AbstractRef,
    AbstractLinVal,
    AccumEffect,
    ReadEffect,
    Transform,
    TransformedRef,
    WriteEffect,
    transform_type as transform_type,
)
from jax._src.typing import Array, ArrayLike
from jax._src.util import safe_map, safe_zip


# Stand-in for hi-jax inputs to Ref.
HijaxType = Any


## General utilities

## JAX utilities

map, unsafe_map = safe_map, map
zip, unsafe_zip = safe_zip, zip
traceback_util.register_exclusion(__file__)

## get/swap/addupdate implementations

# `get` reads a value from a `Ref` type, a.k.a.:
# a = get_p.bind(x)
# or we can read using indices:
# a = get_p.bind(x, 0, 1)
# Staging out `a = get_p.bind(x)` where the aval of `x` is
# `Ref((3,), np.dtype('float32'))` leads to a jaxpr eqn printed like
#   a:f32[3] <- x[]
get_p = core.Primitive("get")
get_p.is_effectful = lambda params: True
get_p.def_impl(partial(dispatch.apply_primitive, get_p))

get_p.is_high = lambda ref_aval, *_, tree: ref_aval.is_high
def _get_to_lojax(ref, *idx, tree):
  val_ty = core.typeof(ref._refs)
  transforms = tree_util.tree_unflatten(tree, idx)
  if transforms:
    ref = TransformedRef(ref, transforms[:-1])
    idx = transforms[-1]
    return val_ty.ref_get_to_lojax(ref, idx)
  return val_ty.raise_val(*map(ref_get, val_ty.lower_val(ref._refs)))
get_p.to_lojax = _get_to_lojax

Indexer = int | slice | Array | types.EllipsisType


def get_ref_and_transforms(
    ref_or_view: Any,
    idx: Indexer | tuple[Indexer, ...] | None,
    function_name: str,
) -> tuple[Any, tuple[Transform, ...]]:
  if isinstance(ref_or_view, TransformedRef):
    ref, transforms = ref_or_view.ref, ref_or_view.transforms
  else:
    ref, transforms = ref_or_view, ()
  ref_aval = core.typeof(ref)
  if not isinstance(ref_aval, AbstractRef):
    raise ValueError(f"Can only call `{function_name}` on a `Ref`: {ref}.")
  if (not isinstance(ref_aval.inner_aval, core.ShapedArray)
      and not ref_aval.inner_aval.is_high):
    return ref, ()

  if idx is None or idx is Ellipsis:
    idx = ()
  elif not isinstance(idx, tuple):
    idx = (idx,)

  if not idx:
    return ref, transforms
  if not idx and transforms and isinstance(transforms[-1], indexing.NDIndexer):
    return ref, transforms
  nd_indexer = indexing.NDIndexer.from_indices_shape(idx, ref_or_view.shape)
  return ref, (*transforms, nd_indexer)

@partial(traceback_util.api_boundary, repro_api_name="jax.ref.get")
def ref_get(
    ref: core.Ref | TransformedRef,
    idx: Indexer | tuple[Indexer, ...] | None = None
) -> Array | HijaxType:
  """Read a value from an Ref.

  This is equivalent to ``ref[idx]`` for a NumPy-style indexer ``idx``.
  For more on mutable array refs, refer to the `Ref guide`_.

  Args:
    ref: a :class:`jax.ref.Ref` object.
    idx: a NumPy-style indexer

  Returns:
    A :class:`jax.Array` object (note, not a :class:`jax.ref.Ref`) containing
    the indexed elements of the mutable reference.

  Examples:
    >>> import jax
    >>> ref = jax.new_ref(jax.numpy.arange(5))
    >>> jax.ref.get(ref, slice(1, 3))
    Array([1, 2], dtype=int32)

    Equivalent operation via indexing syntax:

    >>> ref[1:3]
    Array([1, 2], dtype=int32)

    Use ``...`` to extract the full buffer:

    >>> ref[...]
    Array([0, 1, 2, 3, 4], dtype=int32)

  .. _Ref guide: https://docs.jax.dev/en/latest/101/state.html#jax-101-refs
  """
  if isinstance(ref, TransformedRef) and ref.multiref:
    raise NotImplementedError(
        "ref_get with multi-ref TransformedRef is not supported."
    )
  ref, transforms = get_ref_and_transforms(ref, idx, "ref_get")
  flat_transforms, tree = tree_util.tree_flatten(transforms)
  return get_p.bind(ref, *flat_transforms, tree=tree)


# `swap` mutates a `Ref`, setting its value and returns its previous value.
# b = swap_p.bind(x, a)
# It generalizes the setting operation for a `Ref` as we can ignore the return
# value:
# _ = swap_p.bind(x, a)
# `swap_p` also takes in index arguments following the value, i.e.:
# _ = swap_p.bind(x, a, 0, 1)
# Staging out `b = swap_p.bind(x, a)` where the aval of `x` is
# `Ref((3,), np.dtype('float32'))` and the aval of `a` is
# `ShapedArray((3,), np.dtype('float32'))` leads to a jaxpr eqn printed like
#   b:f32[3], x:Ref{f32[3]} <- x, a
# Staging out `_ = swap_p.bind(x, a, i, j)` where the aval of `x` is
# `Ref((3,), np.dtype('float32'))` , the aval of `a` is
# `ShapedArray((3,), np.dtype('float32'))`, and the avals of both `i` and `j`
# are `ShapedArray((), np.dtype('int32'))` leads to a jaxpr eqn printed like
#   x:Ref{f32[3]}[i, j] <- a
swap_p = core.Primitive("swap")
swap_p.is_effectful = lambda params: True
swap_p.def_impl(partial(dispatch.apply_primitive, swap_p))

swap_p.is_high = lambda ref_aval, *_, tree: ref_aval.is_high
def _swap_to_lojax(ref, val, *idx, tree):
  ref_val_ty = core.typeof(ref._refs)
  val_ty = core.typeof(val)
  transforms = tree_util.tree_unflatten(tree, idx)
  if transforms:
    ref = TransformedRef(ref, transforms[:-1])
    idx = transforms[-1]
    return ref_val_ty.ref_swap_to_lojax(ref, val, idx)
  lo_refs = ref_val_ty.lower_val(ref._refs)
  lo_vals = val_ty.lower_val(val)
  outs = [ref_swap(lo_ref, idx, lo_val) for lo_ref, lo_val
          in zip(lo_refs, lo_vals)]
  return val_ty.raise_val(*outs)
swap_p.to_lojax = _swap_to_lojax


@partial(traceback_util.api_boundary, repro_api_name="jax.ref.swap")
def ref_swap(
    ref: core.Ref | TransformedRef,
    idx: Indexer | tuple[Indexer, ...] | None,
    value: ArrayLike | HijaxType,
    _function_name: str = "ref_swap",
) -> Array | HijaxType:
  """Update an array value inplace while returning the previous value.

  This is equivalent to ``ref[idx], prev = value, ref[idx]`` while returning
  ``prev``, for a NumPy-style indexer ``idx``.
  For more on mutable array refs, refer to the `Ref guide`_.

  Args:
    ref: a :class:`jax.ref.Ref` object. On return, the buffer will be
      mutated by this operation.
    idx: a NumPy-style indexer
    value: a :class:`jax.Array` object (note, not a :class:`jax.ref.Ref`)
      containing the values to set in the array.

  Returns:
    A :class:`jax.Array` containing the previous value at `idx`.

  Examples:
    >>> import jax
    >>> ref = jax.new_ref(jax.numpy.arange(5))
    >>> jax.ref.swap(ref, 3, 10)
    Array(3, dtype=int32)
    >>> ref
    Ref([ 0,  1,  2, 10,  4], dtype=int32)

    Equivalent operation via indexing syntax:

    >>> ref = jax.new_ref(jax.numpy.arange(5))
    >>> ref[3], prev = 10, ref[3]
    >>> prev
    Array(3, dtype=int32)
    >>> ref
    Ref([ 0,  1,  2, 10,  4], dtype=int32)

    Use ``...`` to swap the value of a scalar ref:

    >>> ref = jax.new_ref(jax.numpy.int32(5))
    >>> jax.ref.swap(ref, ..., 10)
    Array(5, dtype=int32)
    >>> ref
    Ref(10, dtype=int32)

  .. _Ref guide: https://docs.jax.dev/en/latest/101/state.html#jax-101-refs
  """
  if isinstance(ref, TransformedRef) and ref.multiref:
    raise NotImplementedError(
        "ref_swap with multi-ref TransformedRef is not supported."
    )
  if hasattr(ref, 'dtype'):
    value = _maybe_implicit_cast(ref.dtype, value)
  ref, transforms = get_ref_and_transforms(ref, idx, _function_name)
  flat_transforms, tree = tree_util.tree_flatten(transforms)
  return swap_p.bind(ref, value, *flat_transforms, tree=tree)

# TODO(slebedev,mattjj): replace with special handling of Python numeric types:
# if (isinstance(value, (int, float, complex)) and
#     value == np.array(value, dtype).item()): return cast
def _maybe_implicit_cast(dtype, value):
  aval = core.typeof(value)
  if not isinstance(aval, core.ShapedArray):
    return value
  if (aval.weak_type and
      (dtypes.issubdtype(dtype, np.floating) and
       dtypes.issubdtype(aval.dtype, np.floating)) or
      (dtypes.issubdtype(dtype, np.integer) and
       dtypes.issubdtype(aval.dtype, np.integer))):
    return lax.convert_element_type(value, dtype)
  return value


@partial(traceback_util.api_boundary, repro_api_name="jax.ref.set")
def ref_set(
    ref: core.Ref | TransformedRef,
    idx: Indexer | tuple[Indexer, ...] | None,
    value: ArrayLike | HijaxType,
) -> None:
  """Set a value in an Ref in-place.

  This is equivalent to ``ref[idx] = value`` for a NumPy-style indexer
  ``idx``. For more on mutable array refs, refer to the `Ref guide`_.

  Args:
    ref: a :class:`jax.ref.Ref` object. On return, the buffer will be
      mutated by this operation.
    idx: a NumPy-style indexer
    value: a :class:`jax.Array` object (note, not a :class:`jax.ref.Ref`)
      containing the values to set in the array.

  Returns:
    None

  Examples:
    >>> import jax
    >>> ref = jax.new_ref(jax.numpy.zeros(5))
    >>> jax.ref.set(ref, 1, 10.0)
    >>> ref
    Ref([ 0., 10.,  0.,  0.,  0.], dtype=float32)

    Equivalent operation via indexing syntax:

    >>> ref = jax.new_ref(jax.numpy.zeros(5))
    >>> ref[1] = 10.0
    >>> ref
    Ref([ 0., 10.,  0.,  0.,  0.], dtype=float32)

    Use ``...`` to set the value of a scalar ref:

    >>> ref = jax.new_ref(jax.numpy.int32(0))
    >>> ref[...] = 4
    >>> ref
    Ref(4, dtype=int32)

  .. _Ref guide: https://docs.jax.dev/en/latest/101/state.html#jax-101-refs
  """
  ref_swap(ref, idx, value, _function_name="ref_set")


# `addupdate_p` mutates a `Ref`, adding a value to its existing value.
# Semantically,
# ```
# addupdate ref a *idx
# ```
# is equivalent to
# ```
# b = get ref *idx
# c = add b x
# _ = swap ref c *idx
# ```
addupdate_p = core.Primitive('addupdate')
addupdate_p.is_effectful = lambda params: True
addupdate_p.multiple_results = True
addupdate_p.def_impl(partial(dispatch.apply_primitive, addupdate_p))


def ref_addupdate(
    ref: core.Ref | TransformedRef,
    idx: Indexer | tuple[Indexer, ...] | None,
    x: ArrayLike | HijaxType,
) -> None:
  """Add to an element in an Ref in-place.

  This is analogous to ``ref[idx] += value`` for a NumPy array ``ref`` and
  NumPy-style indexer ``idx``. However, for an Ref ``ref``, executing
  ``ref[idx] += value`` actually performs a ``ref_get``, add, and ``ref_set``,
  so using this function can be more efficient under autodiff. For more on
  mutable array refs, refer to the `Ref guide`_.

  Args:
    ref: a :class:`jax.ref.Ref` object. On return, the buffer will be
      mutated by this operation.
    idx: a NumPy-style indexer
    x: a :class:`jax.Array` object (note, not a :class:`jax.ref.Ref`)
      containing the values to add at the specified indices.

  Returns:
    None

  Examples:
    >>> import jax
    >>> ref = jax.new_ref(jax.numpy.arange(5))
    >>> jax.ref.addupdate(ref, 2, 10)
    >>> ref
    Ref([ 0,  1, 12,  3,  4], dtype=int32)

    Equivalent operation via indexing syntax:

    >>> ref = jax.new_ref(jax.numpy.arange(5))
    >>> ref[2] += 10
    >>> ref
    Ref([ 0,  1, 12,  3,  4], dtype=int32)

    Use ``...`` to add to a scalar ref:

    >>> ref = jax.new_ref(jax.numpy.int32(2))
    >>> ref[...] += 10
    >>> ref
    Ref(12, dtype=int32)

  .. _Ref guide: https://docs.jax.dev/en/latest/101/state.html#jax-101-refs
  """
  ref, transforms = get_ref_and_transforms(ref, idx, "ref_addupdate")
  flat_transforms, tree = tree_util.tree_flatten(transforms)
  addupdate_p.bind(ref, x, *flat_transforms, tree=tree)


## get/set/addupdate abstract evaluation rules


def _get_abstract_eval(ref_aval: AbstractRef, *args,
                       tree):
  transforms = tree_util.tree_unflatten(tree, args)
  if transforms and ref_aval.inner_aval.is_high:
    # TODO(mattjj): aval.is_high does not imply the existence of ref_get_abstract_aval.
    return ref_aval.inner_aval.ref_get_abstract_eval(ref_aval, *args, tree=tree)  # pyrefly: ignore[missing-attribute]
  if not isinstance(ref_aval, AbstractRef):
    raise ValueError(f"`get` must be called on `Ref` types: {ref_aval}.")
  if isinstance(ref_aval.inner_aval, core.ShapedArray):
    out_aval = transform_type(transforms, ref_aval.inner_aval)
  else:
    if transforms:
      raise ValueError("Cannot index non-shaped array with nontrivial indices.")
    out_aval = ref_aval.inner_aval
  return (out_aval, {ReadEffect(0)})
get_p.def_effectful_abstract_eval(_get_abstract_eval)

def _swap_abstract_eval(ref_aval: AbstractRef,
                        val_aval: core.AbstractValue,
                        *args: Any, tree):
  transforms = tree_util.tree_unflatten(tree, args)
  if transforms and ref_aval.inner_aval.is_high:
    # TODO(mattjj): aval.is_high does not imply the existence of ref_swap_abstract_aval.
    return ref_aval.inner_aval.ref_swap_abstract_eval(  # pyrefly: ignore[missing-attribute]
        ref_aval, val_aval, *args, tree=tree)
  out_aval: core.AbstractValue
  if not isinstance(ref_aval, AbstractRef):
    raise ValueError(f"`swap` must be called on `Ref` types: {ref_aval}.")
  if isinstance(val_aval, AbstractRef):
    raise ValueError("Cannot store a Ref into another Ref. "
                     "Did you forget to load from it using `[...]`?")
  if isinstance(ref_aval.inner_aval, core.ShapedArray):
    assert isinstance(val_aval, core.ShapedArray)
    expected_out_ty = transform_type(transforms, ref_aval.inner_aval)
    assert isinstance(expected_out_ty, core.ShapedArray)
    if expected_out_ty.shape != val_aval.shape:
      raise ValueError("Invalid shape for `swap`. "
                       f"Ref shape: {ref_aval.shape}. "
                       f"Expected shape: {expected_out_ty.shape}. "
                       f"Value shape: {val_aval.shape}. "
                       f"Transforms: {transforms}. ")
    if expected_out_ty.dtype != val_aval.dtype:
      raise ValueError(
          "Invalid dtype for `swap`. "
          f"Ref dtype: {expected_out_ty.dtype}. "
          f"Value dtype: {val_aval.dtype}. "
      )
    out_aval = expected_out_ty
  else:
    if transforms:
      raise ValueError("Cannot index non-shaped array with nontrivial indices.")
    out_aval = ref_aval.inner_aval
  return (out_aval, {WriteEffect(0)})
swap_p.def_effectful_abstract_eval(_swap_abstract_eval)


def _addupdate_abstract_eval(ref_aval: AbstractRef,
                             val_aval: core.AbstractValue,
                             *args: Any, tree):
  transforms = tree_util.tree_unflatten(tree, args)
  if not isinstance(ref_aval, AbstractRef):
    raise ValueError(f"`addupdate` must be called on `Ref` types: {ref_aval}.")
  if isinstance(ref_aval.inner_aval, core.ShapedArray):
    expected_out_ty = transform_type(transforms, ref_aval.inner_aval)
    assert isinstance(val_aval, core.ShapedArray)
    assert isinstance(expected_out_ty, core.ShapedArray)
    if expected_out_ty.shape != val_aval.shape:
      raise ValueError(
          "Invalid shape for `addupdate`. "
          f"Ref shape: {ref_aval.shape}. "
          f"Expected shape: {expected_out_ty.shape}. "
          f"Value shape: {val_aval.shape}. "
          f"Transforms: {transforms}. "
      )
    if expected_out_ty.dtype != val_aval.dtype:
      raise ValueError("Invalid dtype for `addupdate`. "
                       f"Ref dtype: {ref_aval.dtype}. "
                       f"Value shape: {val_aval.dtype}. ")
    out_sharding = expected_out_ty.sharding
    if ((out_sharding.mesh._any_axis_explicit or
         val_aval.sharding.mesh._any_axis_explicit) and
        out_sharding != val_aval.sharding):
      raise ValueError("Invalid sharding for `addupdate`. "
                       f"Ref sharding: {ref_aval.sharding}. "
                       f"Value sharding: {val_aval.sharding}. ")
  else:
    # Check that the transforms are valid
    if transforms:
      raise ValueError("Cannot index non-shaped array with nontrivial indices.")
  return [], {AccumEffect(0)}
addupdate_p.def_effectful_abstract_eval(_addupdate_abstract_eval)

## Pretty printing for `get` and `swap` in jaxprs

pp_ref_var = partial(pp.color, intensity=pp.Intensity.NORMAL,
                 foreground=pp.Color.GREEN)


def _pp_transforms(
    context: core.JaxprPpContext,
    transforms: tuple[Transform, ...],
):
  if not transforms:
    return pp.text("[...]")
  return pp.concat(
      [transform.pretty_print(context) for transform in transforms]
  )


def pp_ref_transforms(context: core.JaxprPpContext, ref, transforms=()):
  if isinstance(ref, TransformedRef):
    transforms = ref.transforms
    ref = ref.ref
  if isinstance(ref, tuple):
    ref_doc = pp.concat([
        pp.text("("),
        pp.join(pp.text(", "), [pp_ref_transforms(context, r) for r in ref]),
        pp.text(")"),
    ])
  else:
    ref_doc = core.pp_var(ref, context)
  return pp_ref_var(
      pp.concat([
          ref_doc,
          _pp_transforms(context, transforms),
      ])
  )


def _get_pp_rule(eqn, context, settings) -> pp.Doc:
  # Pretty prints `a = get x i` as `x[i] <- a`
  y, = eqn.outvars
  x, *flat_idx = eqn.invars
  transforms = tree_util.tree_unflatten(eqn.params["tree"], flat_idx)
  lhs = core.pp_vars([y], context, print_shapes=settings.print_shapes)
  annotation = (source_info_util.summarize(eqn.source_info)
                if settings.source_info else None)
  return pp.concat(
      [lhs, pp.text(" <- ", annotation=annotation),
       pp_ref_transforms(context, x, transforms)]
  )
core.pp_eqn_rules[get_p] = _get_pp_rule

def _swap_pp_rule(eqn, context, settings) -> pp.Doc:
  y, = eqn.outvars
  x, v, *flat_idx = eqn.invars
  transforms = tree_util.tree_unflatten(eqn.params["tree"], flat_idx)
  annotation = (source_info_util.summarize(eqn.source_info)
                if settings.source_info else None)
  if context.var_names.get(y) == '_':
    # In the case of a set (ignored return value), print as `x[i] <- v`
    del y
    return pp.concat([
        pp_ref_transforms(context, x, transforms),
        pp.text(" <- ", annotation=annotation),
        core.pp_var(v, context),
    ])
  else:
    # pretty-print `y:T = swap x v i` as `y:T, x[i] <- x[i], v`
    x_i = pp_ref_transforms(context, x, transforms)
    y = core.pp_vars([y], context, print_shapes=settings.print_shapes)
    return pp.concat([y, pp.text(', '), x_i,
                      pp.text(' <- ', annotation=annotation), x_i,
                      pp.text(', '), core.pp_var(v, context)])
core.pp_eqn_rules[swap_p] = _swap_pp_rule

def _addupdate_pp_rule(eqn, context, settings) -> pp.Doc:
  # pretty-print ` = addupdate x i v` as `x[i] += v`
  () = eqn.outvars
  x, v, *flat_idx = eqn.invars
  transforms = tree_util.tree_unflatten(eqn.params["tree"], flat_idx)
  annotation = (source_info_util.summarize(eqn.source_info)
                if settings.source_info else None)
  return pp.concat([
      pp_ref_transforms(context, x, transforms),
      pp.text(" += ", annotation=annotation),
      core.pp_var(v, context),
  ])
core.pp_eqn_rules[addupdate_p] = _addupdate_pp_rule

## get/swap/addupdate JVP rules

def _get_jvp(primals: list[Any], tangents: list[Any], **params: Any):
  ref_primal, *idx = primals
  ref_tangent, *_ = tangents
  out_primal = get_p.bind(ref_primal, *idx, **params)
  if isinstance(ref_tangent, ad_util.Zero):
    out_tangent = ad_util.Zero(core.typeof(out_primal).to_tangent_aval())
  else:
    out_tangent = get_p.bind(ref_tangent, *idx, **params)
  return out_primal, out_tangent
ad.primitive_jvps[get_p] = _get_jvp

def _swap_jvp(primals: list[Any], tangents: list[Any], **params: Any):
  ref_primal, x_primal, *idx = primals
  ref_tangent, x_tangent, *_ = tangents
  out_primal = swap_p.bind(ref_primal, x_primal, *idx, **params)
  if isinstance(ref_tangent, ad_util.Zero) and isinstance(x_tangent, ad_util.Zero):
    out_tangent = ad_util.Zero(core.typeof(out_primal).to_tangent_aval())
  elif ref_tangent.aval.kind == "no_grad_no_remat":
    out_tangent = ad_util.Zero(core.typeof(out_primal).to_tangent_aval())
  else:
    if isinstance(ref_tangent, ad_util.Zero):
      raise Exception("performing a set/swap operation with a differentiated "
                      "value on a non-differentiated array reference of type "
                      f"{core.typeof(ref_primal)}. Move the array reference "
                      "to be an argument of the differentiated function?")
    x_tangent = ad_util.instantiate(x_tangent)
    out_tangent = swap_p.bind(ref_tangent, x_tangent, *idx, **params)
  return out_primal, out_tangent
ad.primitive_jvps[swap_p] = _swap_jvp

def addupdate_jvp_rule(primals: list[Any], tangents: list[Any], **params: Any):
  ref_primal, x_primal, *idx = primals
  ref_tangent, x_tangent, *_ = tangents
  x_tangent = ad_util.instantiate(x_tangent)
  if ref_tangent.aval.kind != "no_grad_no_remat":
    addupdate_p.bind(ref_primal, x_primal, *idx, **params)
    addupdate_p.bind(ref_tangent, x_tangent, *idx, **params)
  return [], []
ad.primitive_jvps[addupdate_p] = addupdate_jvp_rule

##  get/swap/addupdate transpose rules

def _get_transpose_fancy(g, ref_, *idx, tree):
  transforms = tree_util.tree_unflatten(tree, idx)
  if transforms and type(g) is not ad_util.Zero:
    addupdate_p.bind(ref_.inst().ref, g, *idx, tree=tree)
  else:
    ref_.accum(g)
ad.fancy_transposes[get_p] = _get_transpose_fancy

def _swap_transpose_fancy(g, ref_, x, *idx, **params):
  if ref_.ref is None and type(g) is ad_util.Zero:
    return
  elif ref_.ref is None:
    swap_p.bind(ref_.inst().ref, ad_util.instantiate(g), *idx, **params)
  else:
    x_bar = swap_p.bind(ref_.inst().ref, ad_util.instantiate(g), *idx, **params)
    if isinstance(x, ad.GradAccum):
      x.accum(x_bar)  # if x is a constant, drop x_bar, but still zero the ref ct
ad.fancy_transposes[swap_p] = _swap_transpose_fancy

def addupdate_transpose_fancy(cts_in, ref_, x, *idx, **params):
  if ref_.ref is not None and isinstance(x, ad.GradAccum):
    x_bar = get_p.bind(ref_.ref, *idx, **params)
    x.accum(x_bar)
ad.fancy_transposes[addupdate_p] = addupdate_transpose_fancy

## get/swap/addupdate partial_eval_custom rules

def _array_ref_partial_eval_custom(saveable, unks_in, inst_in, eqn):
  del saveable  # ignored, always full remat array_ref on known input
  unk, = unks_in
  inst, = inst_in
  invar, = eqn.invars
  res = [invar] if not inst else []
  if unk:
    return None, eqn, [True], [True], res  # tangent operation
  else:
    return eqn, eqn, [False], [True], res  # full remat
pe.partial_eval_jaxpr_custom_rules[core.ref_p] = _array_ref_partial_eval_custom

def _array_ref_batched(axis_data, vals_in, dims_in, memory_space, kind, pin):
  val, = vals_in
  dim, = dims_in
  if dim is None:
    # We defensively batch the ref, b/c it could later be hit with a batched val
    val2 = batching.broadcast(val, axis_data.size, 0,
                              axis_data.explicit_mesh_axis)
    return core.ref_p.bind(val2, memory_space=memory_space, kind=kind,
                           pin=pin), 0
  else:
    return core.ref_p.bind(val, memory_space=memory_space, kind=kind,
                           pin=pin), dim
batching.fancy_primitive_batchers[core.ref_p] = _array_ref_batched

def _freeze_batched(axis_data, vals_in, dims_in):
  ref, = vals_in
  dim, = dims_in
  return core.freeze_p.bind(ref), dim
batching.fancy_primitive_batchers[core.freeze_p] = _freeze_batched

def _state_partial_eval_custom(saveable, unks_in, inst_in, eqn):
  del saveable  # ignored, always full remat state ops on known inputs
                # (except for no_grad_no_remat)
  ref_unk, *val_unks = unks_in
  ref_inst, *inst_in = inst_in
  _, *val_vars = eqn.invars
  assert ref_inst
  res = [v for v, inst in zip(val_vars, inst_in) if not inst]
  if ref_unk:
    return None, eqn, [True], [True], res  # tangent operation
  elif any(val_unks):
    raise NotImplementedError(
        "Mutating a reference with unknown/differentiated values is not"
        " supported.\nIf this occurred under `jax.checkpoint` and"
        " `jax.config.jax_remat3 == False`, consider one of the following"
        " options:\n\n"
        "  * Enabling `jax_remat3` via `jax.config.update('jax_remat3',"
        " True)`.\n"
        "  * Creating the Ref outside the function being rematerialized and"
        " passing it as an input.\n"
        "  * Creating the Ref in a data-dependent way from a value being"
        " differentiated, e.g.:\n\n"
        "      ref = jax.new_ref(x * 0.0)  # instead of zeros_like(x)"
    )
  elif eqn.invars[0].aval.kind == "no_grad_no_remat":
    return eqn, None, [False], [False], res
  else:
    return eqn, eqn, [False], [True], res  # full remat
pe.partial_eval_jaxpr_custom_rules[get_p] = _state_partial_eval_custom
pe.partial_eval_jaxpr_custom_rules[swap_p] = _state_partial_eval_custom

def _addupdate_partial_eval_custom(saveable, unks_in, inst_in, eqn):
  del saveable  # ignored, always full remat state ops on known inputs
  ref_unk, *val_unks = unks_in
  ref_inst, *inst_in = inst_in
  _, *val_vars = eqn.invars
  assert ref_inst
  res = [v for v, inst in zip(val_vars, inst_in) if not inst]
  if ref_unk:
    return None, eqn, [], [], res  # tangent operation
  elif any(val_unks):
    raise NotImplementedError(
        "Mutating a reference with unknown/differentiated values is not"
        " supported.\nIf this occurred under `jax.checkpoint` and"
        " `jax.config.jax_remat3 == False`, consider one of the following"
        " options:\n\n"
        "  * Enabling `jax_remat3` via `jax.config.update('jax_remat3',"
        " True)`.\n"
        "  * Creating the Ref outside the function being rematerialized and"
        " passing it as an input.\n"
        "  * Creating the Ref in a data-dependent way from a value being"
        " differentiated, e.g.:\n\n"
        "      ref = jax.new_ref(x * 0.0)  # instead of zeros_like(x)"
    )
  else:
    return eqn, eqn, [], [], res  # full remat
pe.partial_eval_jaxpr_custom_rules[addupdate_p] = _addupdate_partial_eval_custom

##  get/swap/addupdate remat rules

def _state_remat(prim, trace, ref, *args, **params):
  del trace
  out = prim.bind(ref, *args, **params)
  if core.typeof(ref).kind == "no_grad_no_remat":
    # Run only in the fwd computation; any read value reaching the remnant
    # computation is saved as a residual rather than re-read.
    return out, out, lambda out, *_: out
  return out, (), lambda _, ref_, *args_: prim.bind(ref_, *args_, **params)

remat.rules[get_p] = partial(_state_remat, get_p)
remat.rules[swap_p] = partial(_state_remat, swap_p)
remat.rules[addupdate_p] = partial(_state_remat, addupdate_p)

##  get/swap/addupdate batching rules

def _batch_indexer(
    indexer: indexing.NDIndexer,
    dims,
    axis_size: int,
    ref_shape: tuple[int, ...],
    ref_dim: int | batching.NotMapped,
    idx_is_batched: bool,
) -> indexing.NDIndexer:
  """Converts a batched indexer into an unbatched one.

  This function handles the complexity of `vmap`-style batching where either the
  `ref` being indexed, the indexer, or both may have batched dimensions. The
  goal is to produce a new indexer that acts as if applied in a batched context,
  but without actual batching, enabling downstream code to process it as usual.

  If any index in `indexer` is batched, all array indexers are normalized. If
  the array indexer contains a batched dimension, the dimension is moved to the
  front (axis 0). If the array indexer not batched, it is broadcasted to include
  a batch dimension at the front. This is to guarantee that all array indexers
  are still of the same shape.

  Slices are passed through unchanged unless they contain dynamic elements and
  are themselves batched, which is currently unsupported.

  If `ref` is batched (`ref_dim` is not `NotMapped`), we simulate per-example
  indexing by inserting a new iota array at the position corresponding to
  `ref_dim` in the indexer.

  It is worth noting that if the array indexers in the original indexer are
  contiguous, but become non-contiguous in the new indexer due to the insertion
  of the iota, the dimensions corresponding to the array indexers will be moved
  to the front in the indexing result. The batched dimension will be at axis 0,
  while the dimensions corresponding to the array indexers in the original
  indexer will start from axis 1. This behavior would cause a mismatch between
  the original indexer and the new indexer. Callers must take this behavior into
  account and properly transpose the arrays involved to avoid this mismatch.

  Args:
    indexer: An `NDIndexer` that indexes into `ref`.
    dims: A pytree with the same structure as `indexer`, indicating which
      dimension (if any) is batched for each array indexer.
    axis_size: Size of the batch dimension.
    ref_shape: Shape of `ref`.
    ref_dim: The dimension of `ref` that is batched (if any).
    idx_is_batched: Whether any index in the `indexer` is batched.
  """
  indices = indexer.indices
  indices_dims = dims.indices
  new_indices: list[Array | indexing.Slice | int] = []
  new_integer_indexer_shape = (axis_size, *indexer.int_indexer_shape)
  for idx, dim in zip(indices, indices_dims):
    if idx_is_batched:
      # If at least one of the idx is batched, we broadcast them all and move the
      # batch dim to the front.
      if isinstance(idx, indexing.Slice):
        # size is static, but start can be dynamic
        # Check if start is static (which it can be)
        is_static_slice = len(tree_util.tree_leaves(idx)) == 0
        if is_static_slice:
          new_indices.append(idx)
          continue
        dim = dim.start
        if dim is None:
          # Broadcasting the slice is free (the start index stays the same)
          new_indices.append(idx)
        else:
          raise NotImplementedError(
              f"No support for vmapping over nontrivial slices just yet: {idx}")
      else:
        # Check if we are indexing with a scalar or not. If we are indexing
        # with a scalar and we are not batched, we can avoid broadcasting it.
        if not shapeof(idx):
          new_indices.append(idx)
        else:
          if dim is None:
            bcast_dims = tuple(range(1, np.ndim(idx) + 1))
            idx = lax.broadcast_in_dim(idx, new_integer_indexer_shape,
                                       bcast_dims)
          else:
            idx = batching.moveaxis(idx, dim, 0)
          new_indices.append(idx)
    else:
      if ref_dim is not None:
        if not isinstance(idx, indexing.Slice):
          if shapeof(idx):
            bcast_dims = tuple(range(1, core.typeof(idx).ndim + 1))
            idx = lax.broadcast_in_dim(idx, new_integer_indexer_shape,
                                       bcast_dims)
      new_indices.append(idx)

  # if the ref is batched, we must insert an indexer for that new axis
  if ref_dim is not None:
    # sometimes can optimize to generate a slice(None) on the batch axis.
    # NOTE this optimization has caused much pain. if in doubt, try removing.
    if not idx_is_batched and not indexer.int_indexer_shape:
      batch_idx = indexing.Slice(0, axis_size)
      new_integer_indexer_shape = ()
    else:
      batch_idx = lax.broadcasted_iota(
          np.dtype('int32'), new_integer_indexer_shape, 0)
    new_indices.insert(ref_dim, batch_idx)

  return indexing.NDIndexer(
      tuple(new_indices), ref_shape, new_integer_indexer_shape, validate=True
  )

def shapeof(x):
  return x.shape if isinstance(x, TransformedRef) else core.typeof(x).shape

def _get_vmap(batched_args, batched_dims, *, tree):
  axis_size, = {x.shape[d] for x, d in zip(batched_args, batched_dims)
                if d is not None}
  ref, *flat_idxs = batched_args
  ref_dim, *flat_idx_dims = batched_dims
  indexers = tree_util.tree_unflatten(tree, flat_idxs)
  if not indexers:
    return get_p.bind(ref, *flat_idxs, tree=tree), ref_dim
  indexers_dims = tree_util.tree_unflatten(tree, flat_idx_dims)

  idx_is_batched = any(i_dim is not None
                       for i_dim in flat_idx_dims)
  if len(indexers) > 1:
    raise NotImplementedError("Batching with multiple indexers not supported.")

  # TODO(sharadmv): handle vmap of multiple indexers
  new_indexers = tuple(_batch_indexer(indexer, dims, axis_size,
                                  ref.shape, ref_dim, idx_is_batched)
                     for indexer, dims in zip(indexers, indexers_dims))
  flat_indexers, tree = tree_util.tree_flatten(new_indexers)

  is_int_indexing, _, _ = indexing.unpack_ndindexer(indexers[0])
  int_indexers_contiguous = bool(
      np.all(np.diff(np.where(is_int_indexing)[0]) == 1)
  )
  # Note: _batch_indexer will add a slice for the batch dim if the int_indexer
  # shape is empty, else it will use advanced/int indexing.
  will_add_int_batcher = ref_dim is not None and (idx_is_batched or indexers[0].int_indexer_shape)

  is_new_int_indexing, _, _ = indexing.unpack_ndindexer(new_indexers[0])
  new_int_indexers_contiguous = bool(
      np.all(np.diff(np.where(is_new_int_indexing)[0]) == 1)
  )

  out = get_p.bind(ref, *flat_indexers, tree=tree)
  should_transpose = (int_indexers_contiguous and
                      not new_int_indexers_contiguous)
  if will_add_int_batcher and should_transpose:
    original_pos = is_int_indexing.index(True)
    array_indexer_shape = new_indexers[0].int_indexer_shape
    array_indexer_len = len(array_indexer_shape)

    transpose_order = list(range(len(out.shape)))
    transpose_order = (
        transpose_order[0],
        *transpose_order[array_indexer_len:array_indexer_len+original_pos],
        *transpose_order[1:array_indexer_len],
        *transpose_order[array_indexer_len+original_pos:],
    )
    out = lax.transpose(out, transpose_order)
    out_bdim = 0
  else:
    if ref_dim is not None:
      if will_add_int_batcher:
        if not int_indexers_contiguous:
          # In this case the indexer is always moved to the front.
          out_bdim = 0
        else:
          # In this case the indexer is not moved to the front.
          out_bdim = is_new_int_indexing.index(True)
      else:
        # We only trigger this case when the int_indexer shape is empty,
        # so we don't need to account for int_indexer_shape.
        int_indexers_before_ref_dim = sum(is_new_int_indexing[:ref_dim])
        out_bdim = ref_dim - int_indexers_before_ref_dim
    else:
      out_bdim = 0
      if new_int_indexers_contiguous:
        out_bdim = is_new_int_indexing.index(True)
  return out, out_bdim
batching.primitive_batchers[get_p] = _get_vmap

def _swap_vmap(axis_data, batched_args, batched_dims, *, tree):
  ref, val, *flat_idxs = batched_args
  ref_dim, val_dim, *flat_idx_dims = batched_dims
  indexers = tree_util.tree_unflatten(tree, flat_idxs)
  indexers_dims = tree_util.tree_unflatten(tree, flat_idx_dims)

  ref_is_batched = ref_dim is not None
  val_is_batched = val_dim is not None
  idx_is_batched = any(i_dim is not None
                       for i_dim in flat_idx_dims)

  if not ref_is_batched:
    raise Exception("performing a set/swap operation with vmapped value on "
                    f"an unbatched array reference of type {core.typeof(ref)}. "
                    "Move the array reference to be an argument to the vmapped "
                    "function?")
  if not indexers:
    if ref_is_batched and not val_is_batched:
      val = batching.broadcast(val, axis_data.size, ref_dim,
                               axis_data.explicit_mesh_axis)
    else:
      val = batching.moveaxis(val, val_dim, ref_dim)
    return swap_p.bind(ref, val, *flat_idxs, tree=tree), ref_dim
  if len(indexers) > 1:
    raise NotImplementedError("Batching with multiple indexers not supported.")
  # TODO(sharadmv): handle vmap of multiple indexers
  new_indexers = tuple(_batch_indexer(indexer, dims, axis_data.size,
                                      ref.shape, ref_dim, idx_is_batched)
                       for indexer, dims in zip(indexers, indexers_dims))
  flat_indexers, tree = tree_util.tree_flatten(new_indexers)

  is_int_indexing, _, _ = indexing.unpack_ndindexer(indexers[0])
  int_indexers_contiguous = bool(
      np.all(np.diff(np.where(is_int_indexing)[0]) == 1)
  )
  is_new_int_indexing, _, _ = indexing.unpack_ndindexer(new_indexers[0])
  new_int_indexers_contiguous = bool(
      np.all(np.diff(np.where(is_new_int_indexing)[0]) == 1)
  )

  # Note: _batch_indexer will add a slice for the batch dim if the int_indexer
  # shape is empty, else it will use advanced/int indexing.
  will_add_int_batcher = idx_is_batched or indexers[0].int_indexer_shape

  if not new_int_indexers_contiguous and will_add_int_batcher:
    batched_dim_in_result = 0  # will be moved to the front
  elif any(is_new_int_indexing):
    if will_add_int_batcher:
      # _batch_indexer ensures the bdim in the int indexers is at the front
      batched_dim_in_result = is_new_int_indexing.index(True)
    else:
      # all int indices are scalars, so the ones before ref_dim disappear
      assert ref_dim is not None
      batched_dim_in_result = ref_dim - sum(is_new_int_indexing[:ref_dim])
  else:
    batched_dim_in_result = ref_dim

  if not val_is_batched:
    if ref_is_batched or idx_is_batched:
      val = batching.broadcast(val, axis_data.size, batched_dim_in_result,
                               axis_data.explicit_mesh_axis)
  else:
    assert batched_dim_in_result is not None
    val = batching.moveaxis(val, val_dim, batched_dim_in_result)

  # Originally not going to be moved to front, but now will be moved to front.
  transpose_order_inversed = None
  if int_indexers_contiguous and not new_int_indexers_contiguous and will_add_int_batcher:
    original_pos = is_int_indexing.index(True)
    array_indexer_shape = new_indexers[0].int_indexer_shape
    array_indexer_len = len(array_indexer_shape)

    transpose_order = list(range(len(val.shape)))
    transpose_order = (
        transpose_order[0],
        *transpose_order[1+original_pos:(1+original_pos)+(array_indexer_len-1)],
        *transpose_order[1:1+original_pos],
        *transpose_order[(1+original_pos)+(array_indexer_len-1):],
    )
    val = val.transpose(transpose_order)
    transpose_order_inversed = np.argsort(transpose_order)

  out = swap_p.bind(ref, val, *flat_indexers, tree=tree)

  # `val` should not be transposed, but we needed to transpose it to match
  # `swap_p`. As a result, the output of `swap_p` is also transposed. Now we
  # need to transpose it back.
  if transpose_order_inversed is not None:
    out = out.transpose(transpose_order_inversed)

  return out, batched_dim_in_result
batching.fancy_primitive_batchers[swap_p] = _swap_vmap

def _addupdate_vmap(axis_data, batched_args, batched_dims, *, tree):
  ref, val, *flat_idxs = batched_args
  ref_dim, val_dim, *flat_idx_dims = batched_dims
  indexers = tree_util.tree_unflatten(tree, flat_idxs)
  indexers_dims = tree_util.tree_unflatten(tree, flat_idx_dims)

  ref_is_batched = ref_dim is not None
  val_is_batched = val_dim is not None
  idx_is_batched = any(i_dim is not None
                       for i_dim in flat_idx_dims)

  if not ref_is_batched:
    raise Exception("performing an addupdate operation with vmapped value on "
                    f"an unbatched array reference of type {core.typeof(ref)}. "
                    "Move the array reference to be an argument to the vmapped "
                    "function?")
  if not indexers:
    if val_dim != ref_dim:
      val = batching.matchaxis(axis_data, val_dim, ref_dim, val)
    return addupdate_p.bind(ref, val, *flat_idxs, tree=tree), []
  if len(indexers) > 1:
    raise NotImplementedError("Batching with multiple indexers not supported.")

  # TODO(sharadmv): handle vmap of multiple indexers
  new_indexers = tuple(_batch_indexer(indexer, dims, axis_data.size,
                                      ref.shape, ref_dim, idx_is_batched)
                       for indexer, dims in zip(indexers, indexers_dims))
  flat_indexers, tree = tree_util.tree_flatten(new_indexers)

  is_int_indexing, _, _ = indexing.unpack_ndindexer(indexers[0])
  int_indexers_contiguous = bool(
      np.all(np.diff(np.where(is_int_indexing)[0]) == 1)
  )
  is_new_int_indexing, _, _ = indexing.unpack_ndindexer(new_indexers[0])
  new_int_indexers_contiguous = bool(
      np.all(np.diff(np.where(is_new_int_indexing)[0]) == 1)
  )

  # Note: _batch_indexer will add a slice for the batch dim if the int_indexer
  # shape is empty, else it will use advanced/int indexing.
  will_add_int_batcher = idx_is_batched or indexers[0].int_indexer_shape

  if not new_int_indexers_contiguous and will_add_int_batcher:
    batched_dim_in_result = 0  # will be moved to the front
  elif any(is_new_int_indexing):
    if will_add_int_batcher:
      # _batch_indexer ensures the bdim in the int indexers is at the front
      batched_dim_in_result = is_new_int_indexing.index(True)
    else:
      # all int indices are scalars, so the ones before ref_dim disappear
      assert ref_dim is not None
      batched_dim_in_result = ref_dim - sum(is_new_int_indexing[:ref_dim])
  else:
    batched_dim_in_result = ref_dim

  if not val_is_batched:
    if ref_is_batched or idx_is_batched:
      val = batching.broadcast(val, axis_data.size, batched_dim_in_result,
                               axis_data.explicit_mesh_axis)
  else:
    assert batched_dim_in_result is not None
    val = batching.moveaxis(val, val_dim, batched_dim_in_result)

  # Originally not going to be moved to the front, but now going to be moved to
  # the front.
  if int_indexers_contiguous and not new_int_indexers_contiguous and will_add_int_batcher:
    original_pos = is_int_indexing.index(True)
    array_indexer_shape = new_indexers[0].int_indexer_shape
    array_indexer_len = len(array_indexer_shape)

    transpose_order = list(range(len(val.shape)))
    transpose_order = (
        transpose_order[0],
        *transpose_order[1+original_pos:(1+original_pos)+(array_indexer_len-1)],
        *transpose_order[1:1+original_pos],
        *transpose_order[(1+original_pos)+(array_indexer_len-1):],
    )
    val = val.transpose(transpose_order)

  return addupdate_p.bind(ref, val, *flat_indexers, tree=tree), []
batching.fancy_primitive_batchers[addupdate_p] = _addupdate_vmap

# Currently, JAX doesn't have a primitive that does an equal-rank broadcast.
# We could use `jnp.broadcast_to` but that lowers to squeezing,
# then broadcast_in_dim. Triton has an equal-rank broadcast (`tl.broadcast_to`)
# so in the lowering, we have to expand out those squeezed dimensions again.
# Having a simple `broadcast_to` primitive allows us to lower directly
# to `tl.broadcast_to`.
broadcast_to_p = core.Primitive('broadcast_to')

def broadcast_to(a: Array, shape: tuple[int, ...]) -> Array:
  """Broadcasts an array to a new shape.

  Args:
    a: The array to broadcast.
    shape: The desired shape to broadcast to.

  Returns:
    An array of shape ``shape``.

  See Also:
    :func:`jax.numpy.broadcast_to`
  """
  import jax.numpy as jnp  # pyrefly: ignore[missing-import]
  a = jnp.asarray(a)
  if a.shape == shape:
    return a
  return broadcast_to_p.bind(a, shape=shape)

@broadcast_to_p.def_impl
def _broadcast_to_impl(a, *, shape):
  import jax.numpy as jnp  # pyrefly: ignore[missing-import]
  return jnp.broadcast_to(a, shape)

@broadcast_to_p.def_abstract_eval
def _broadcast_to_abstract_eval(aval, *, shape):
  return core.ShapedArray(shape, aval.dtype)

mlir.register_lowering(
    broadcast_to_p, mlir.lower_fun(_broadcast_to_impl, False)
)

# === AD rules for mutable arrays ===

def _ref_jvp(primals, tangents, *, memory_space, kind, pin):
  (init_val,), (init_dot,) = primals, tangents
  primal_out = core.ref_p.bind(init_val, memory_space=memory_space, kind=kind,
                               pin=pin)
  if type(init_dot) is ad_util.Zero:
    zero = ad_util.zeros_like_aval(init_dot.aval)
    tangent_out = core.ref_p.bind(zero, memory_space=memory_space, kind=kind,
                                  pin=pin)
  else:
    tangent_out = core.ref_p.bind(init_dot, memory_space=memory_space,
                                  kind=kind, pin=pin)
  return primal_out, tangent_out

def _ref_lin(_is_vjp, nzs, x, *, memory_space, kind, pin):
  nz, = nzs
  x_ref = core.ref_p.bind(x, memory_space=memory_space, kind=kind, pin=pin)
  def mut_lin(_, __, x_dot):
    if kind == 'no_grad_no_remat':
      aval = x_dot.aval if type(x_dot) is ad.Zero else core.typeof(x_dot)
      return ad.Zero(AbstractRef(aval))
    zero = ad_util.instantiate(x_dot)
    return core.ref_p.bind(zero, memory_space=memory_space, kind=kind, pin=pin)
  return x_ref, kind != 'no_grad_no_remat', None, None, mut_lin

ad.primitive_jvps[core.ref_p] = _ref_jvp
ad.primitive_linearizations[core.ref_p] = _ref_lin
ad.linearize_on_zero_tangents.add(core.ref_p)
# TODO(mattjj): lin rule for freeze and accum_grad_in_ref?
ad.defjvp(core.freeze_p, lambda g, _: core.freeze(g))
ad.defjvp(core.accum_grad_in_ref_p, lambda g, _: core.accum_grad_in_ref_p.bind(g))


def _empty_ref_jvp(primals, tangents, *, ty, memory_space, pin):
  primal_ref = core.empty_ref_p.bind(ty=ty, memory_space=memory_space, pin=pin)
  tangent_ref = core.empty_ref_p.bind(ty=ty.to_tangent_aval(),
                                      memory_space=memory_space, pin=pin)
  return primal_ref, tangent_ref
ad.primitive_jvps[core.empty_ref_p] = _empty_ref_jvp

def _empty_ref_lin(_is_vjp, nzs_in, *, ty, memory_space, pin):
  primal_ref = core.empty_ref_p.bind(ty=ty, memory_space=memory_space, pin=pin)
  def lin(_, __):
    return core.empty_ref_p.bind(ty=ty.to_tangent_aval(),
                                 memory_space=memory_space, pin=pin)
  return primal_ref, True, None, None, lin
ad.primitive_linearizations[core.empty_ref_p] = _empty_ref_lin
ad.linearize_on_zero_tangents.add(core.empty_ref_p)

def _free_ref_jvp(primals, tangents):
  [primal_ref], [tangent_ref] = primals, tangents
  core.free_ref(primal_ref)
  core.free_ref(tangent_ref)
  return (), ()

ad.primitive_jvps[core.free_ref_p] = _free_ref_jvp

# === pinned, chained LinearVals ===

def create_linear(ty, memory_space=None):
  return create_linear_p.bind(ty=ty, memory_space=memory_space)
create_linear_p = core.Primitive('create_linear')

@create_linear_p.def_abstract_eval
def _create_linear_abstract_eval(*, ty, memory_space):
  if not isinstance(ty, core.ShapedArray): raise NotImplementedError(ty)
  return AbstractLinVal(ty, memory_space)

def _lower_create_linear(ctx, *, ty, memory_space):
  del ty, memory_space
  out_aval, = ctx.avals_out
  flat_res_types, _ = mlir.ir_tree_registry.flatten(mlir.aval_to_ir_types(ctx.module_context, out_aval))
  return mlir.custom_call(
      "CreateBuffer",
      operands=[],
      result_types=flat_res_types,
  ).results
mlir.register_lowering(create_linear_p, _lower_create_linear)


def pin(x, *, to=None):
  return pin_p.bind(x, to=to)
pin_p = core.Primitive('pin')

@pin_p.def_abstract_eval
def _pin_abstract_eval(aval, *, to):
  if to not in (None, 'hbm', 'vmem'): raise ValueError
  if not isinstance(aval, core.ShapedArray): raise NotImplementedError(aval)
  return AbstractLinVal(aval, to)

def _lower_pin(ctx, x_op, *, to):
  color = {'vmem': 1, 'hbm': 0, None: None}[to]
  if color is not None:
    backend_config = json.dumps({
        "custom_call_config": {
            "output_memory_space_colors": [
                {
                    "shape_index": [],
                    "color": str(color)
                }
            ]
        }
    })
    config: dict[str, Any] = dict(backend_config=backend_config)
  else:
    config = {}
  out_aval, = ctx.avals_out
  flat_ops, _ = mlir.ir_tree_registry.flatten([x_op])
  flat_res_types, _ = mlir.ir_tree_registry.flatten(mlir.aval_to_ir_types(ctx.module_context, out_aval))
  return mlir.custom_call(
      "Pin",
      operands=flat_ops,
      result_types=flat_res_types,
      **config,
  ).results
mlir.register_lowering(pin_p, _lower_pin)


def unpin(x):
  return unpin_p.bind(x)
unpin_p = core.Primitive('unpin')

@unpin_p.def_abstract_eval
def _unpin_abstract_eval(aval):
  if not isinstance(aval, AbstractLinVal): raise TypeError(aval)
  return aval.inner_aval

def _lower_unpin(ctx, x_op):
  out_aval, = ctx.avals_out
  flat_ops, _ = mlir.ir_tree_registry.flatten([x_op])
  flat_res_types, _ = mlir.ir_tree_registry.flatten(mlir.aval_to_ir_types(ctx.module_context, out_aval))
  return mlir.custom_call(
      "Unpin",
      operands=flat_ops,
      result_types=flat_res_types,
  ).results
mlir.register_lowering(unpin_p, _lower_unpin)


def _linval_to_mlir_type(a):
  color = {'vmem': 1, 'hbm': 0, None: None}[a.memory_space]
  space = mlir.i32_attr(color) if color is not None else None
  return mlir.ir.MemRefType.get(a.shape, mlir.dtype_to_ir_type(a.dtype),
                                memory_space=space)
mlir.ir_type_handlers[AbstractLinVal] = _linval_to_mlir_type


with_memory_space_constraint_p = core.Primitive(
    'with_memory_space_constraint'
)


@with_memory_space_constraint_p.def_impl
def _with_memory_space_constraint_impl(x, *, memory_space):
  del memory_space
  return x


@with_memory_space_constraint_p.def_abstract_eval
def _with_memory_space_constraint_abstract_eval(x, *, memory_space):
  if not isinstance(x, core.ShapedArray):
    raise NotImplementedError(
        "with_memory_space_constraint only supports arrays."
    )
  return x.update(memory_space=memory_space)


def _with_memory_space_constraint_lowering_rule(ctx, x, *, memory_space):
  del ctx, memory_space
  return [x]


mlir.register_lowering(
    with_memory_space_constraint_p, _with_memory_space_constraint_lowering_rule
)
batching.defvectorized(with_memory_space_constraint_p)
ad.deflinear2(
    with_memory_space_constraint_p,
    lambda ct, _, **params: (
        with_memory_space_constraint_p.bind(ct, **params),
    ),
)


def with_memory_space_constraint(
    x: Array, memory_space: Any
) -> Array:
  """Constrains the memory space of an array.

  This primitive does not change the value of ``x``, but it constrains the
  memory space where it should be allocated. This is useful to force
  Pallas to allocate an array in a specific memory space.

  As of now, this only operates on the inputs pallas_calls, as in you can
  apply this to the arguments of a pallas_call and it will constrain them, but
  not on intermediate values inside a pallas_call or values after a pallas_call.

  Args:
    x: An array.
    memory_space: The memory space to constrain ``x`` to.

  Returns:
    An array with the same value as ``x``, but with the memory space
    constrained.
  """
  return with_memory_space_constraint_p.bind(x, memory_space=memory_space)