class_docstr: Functions that operate element-by-element on whole arrays.
  
  This is a class for LAX-backed implementations of numpy ufuncs.
code: def _at_via_scan(
      self, a: ArrayLike, indices: Any, *args: Any
  ) -> Array:
      assert len(args) in {0, 1}
      check_arraylike(f"{self.__name__}.at", a, *args)
      dtype = jax.eval_shape(
          self._func,
          lax_internal._one(a),
          *(lax_internal._one(arg) for arg in args),
      ).dtype
      a = lax_internal.asarray(a).astype(dtype)
      args = tuple(
          lax_internal.asarray(arg).astype(dtype) for arg in args
      )
      indices = _eliminate_deprecated_list_indexing(indices)
      if not indices:
          return a
  
      shapes = [
          np.shape(i) for i in indices if not isinstance(i, slice)
      ]
      shape = shapes and jax.lax.broadcast_shapes(*shapes)
      if not shape:
          return a.at[indices].set(
              self._call(a.at[indices].get(), *args)
          )
  
      if args:
          arg = _broadcast_to(
              args[0], (*shape, *args[0].shape[len(shape) :])
          )
          args = (
              arg.reshape(
                  math.prod(shape), *args[0].shape[len(shape) :]
              ),
          )
      indices = [
          idx
          if isinstance(idx, slice)
          else _broadcast_to(idx, shape).ravel()
          for idx in indices
      ]
  
      def scan_fun(carry, x):
          i, a = carry
          idx = tuple(
              ind if isinstance(ind, slice) else ind[i]
              for ind in indices
          )
          a = a.at[idx].set(
              self._call(a.at[idx].get(), *(arg[i] for arg in args))
          )
          return (i + 1, a), x
  
      carry, _ = jax.lax.scan(scan_fun, (0, a), None, len(indices[0]))
      return carry[1]
  
imports: from __future__ import annotations
  from functools import partial
  import math
  import operator
  from typing import Any, Callable
  import jax
  from jax._src.typing import Array, ArrayLike, DTypeLike
  from jax._src.lax import lax
  from jax._src.numpy import reductions
  from jax._src.numpy.lax_numpy import _eliminate_deprecated_list_indexing, append, take
  from jax._src.numpy.reductions import _moveaxis
  from jax._src.numpy.util import _wraps, check_arraylike, _broadcast_to, _where
  from jax._src.numpy.vectorize import vectorize
  from jax._src.util import canonicalize_axis, set_module
  import numpy
label_desc: Use image node_id 11 for calling the ufunc obj's underlying member method code with example usage: obj._at_via_scan(a, indices) and returns: carry, a, a, unknown,  a,  x