Replace all deprecated `jax.ops` operations with jnp's `at`
Author: sanchit-gandhi
Pull request: yes

The `jax.ops.index_update` function used for in-place `jax.ndarray` operations was deprecated in [JAX 0.2.22](https://jax.readthedocs.io/en/latest/changelog.html?highlight=0.2.22#jax-0-2-22-oct-12-2021). Following the official advice, this PR updates any in-place array modifications to be made with [`jax.numpy.ndarray.at`](https://jax.readthedocs.io/en/latest/_autosummary/jax.numpy.ndarray.at.html), thus addressing the issue raised in #16022. Note that in [JAX 0.3.2 (unreleased)](https://jax.readthedocs.io/en/latest/changelog.html#jax-0-3-2-unreleased), this function is removed entirely.

Comments:
_The documentation is not available anymore as the PR was closed or merged._

Cool, thanks for the clean-up!

@patil-suraj - could you also take a look here?

Hi! I am still getting the error that was fixed in this issue, and in #16401, despite confirming that I am running the most up to date version of the repo. I am getting the error when I try to load in a model that uses the transformers/models/bart/**modeling_flax_bart.py** file. Below is the full contents of the error:


---------------------------------------------------------------------------
AttributeError                            Traceback (most recent call last)
[<ipython-input-22-fbba08678499>](https://localhost:8080/#) in <module>()
      6 # Load dalle-mini
      7 model = DalleBart.from_pretrained(
----> 8     DALLE_MODEL, revision=DALLE_COMMIT_ID, dtype=dtype, abstract_init=True
      9 )
     10 

11 frames
[/usr/local/lib/python3.7/dist-packages/dalle_mini/model/utils.py](https://localhost:8080/#) in from_pretrained(cls, pretrained_model_name_or_path, *model_args, **kwargs)
     24 
     25             return super(PretrainedFromWandbMixin, cls).from_pretrained(
---> 26                 pretrained_model_name_or_path, *model_args, **kwargs
     27             )

[/usr/local/lib/python3.7/dist-packages/dalle_mini/model/modeling.py](https://localhost:8080/#) in from_pretrained(cls, pretrained_model_name_or_path, dtype, *model_args, **kwargs)
   1159 
   1160         # init random models
-> 1161         model = cls(config, *model_args, **model_kwargs)
   1162 
   1163         with open(resolved_archive_file, "rb") as state_f:

[/usr/local/lib/python3.7/dist-packages/dalle_mini/model/modeling.py](https://localhost:8080/#) in __init__(self, config, input_shape, seed, dtype, abstract_init, load_on_cpu, init_weights, **kwargs)
   1006                 input_shape,
   1007                 abstract_init=abstract_init,
-> 1008                 load_on_cpu=load_on_cpu,
   1009             )
   1010 

[/usr/local/lib/python3.7/dist-packages/dalle_mini/model/modeling.py](https://localhost:8080/#) in init_weights(self, rng, input_shape, abstract_init, load_on_cpu)
   1024             # only set shape and dtype, load parameters separately
   1025             init_fn = partial(init_fn, input_shape=input_shape)
-> 1026             params = jax.eval_shape(init_fn, rng)
   1027         else:
   1028             params = init_fn(rng, input_shape)

[/usr/local/lib/python3.7/dist-packages/jax/_src/api.py](https://localhost:8080/#) in eval_shape(fun, *args, **kwargs)
   2927   out = pe.abstract_eval_fun(wrapped_fun.call_wrapped,
   2928                              *map(shaped_abstractify, args_flat),
-> 2929                              debug_info=debug_info)
   2930   out = [ShapeDtypeStruct(x.shape, x.dtype, x.named_shape) for x in out]
   2931   return tree_unflatten(out_tree(), out)

[/usr/local/lib/python3.7/dist-packages/jax/interpreters/partial_eval.py](https://localhost:8080/#) in abstract_eval_fun(fun, debug_info, *avals, **params)
    513 def abstract_eval_fun(fun, *avals, debug_info=None, **params):
    514   _, avals_out, _ = trace_to_jaxpr_dynamic(
--> 515       lu.wrap_init(fun, params), avals, debug_info)
    516   assert all(isinstance(aval, AbstractValue) for aval in avals_out)
    517   return avals_out

[/usr/local/lib/python3.7/dist-packages/jax/_src/profiler.py](https://localhost:8080/#) in wrapper(*args, **kwargs)
    204   def wrapper(*args, **kwargs):
    205     with TraceAnnotation(name, **decorator_kwargs):
--> 206       return func(*args, **kwargs)
    207     return wrapper
    208   return wrapper

[/usr/local/lib/python3.7/dist-packages/jax/interpreters/partial_eval.py](https://localhost:8080/#) in trace_to_jaxpr_dynamic(fun, in_avals, debug_info, keep_inputs)
   1737     main.jaxpr_stack = ()  # type: ignore
   1738     jaxpr, out_avals, consts = trace_to_subjaxpr_dynamic(
-> 1739       fun, main, in_avals, keep_inputs=keep_inputs)
   1740     del main, fun
   1741   return jaxpr, out_avals, consts

[/usr/local/lib/python3.7/dist-packages/jax/interpreters/partial_eval.py](https://localhost:8080/#) in trace_to_subjaxpr_dynamic(fun, main, in_avals, keep_inputs)
   1773     in_tracers = _avals_to_tracers(trace, in_avals)
   1774     in_tracers_ = [t for t, keep in zip(in_tracers, keep_inputs) if keep]
-> 1775     ans = fun.call_wrapped(*in_tracers_)
   1776     out_tracers = map(trace.full_raise, ans)
   1777     jaxpr, consts = frame.to_jaxpr(out_tracers)

[/usr/local/lib/python3.7/dist-packages/jax/linear_util.py](https://localhost:8080/#) in call_wrapped(self, *args, **kwargs)
    164 
    165     try:
--> 166       ans = self.f(*args, **dict(self.params, **kwargs))
    167     except:
    168       # Some transformations yield from inside context managers, so we have to

[/usr/local/lib/python3.7/dist-packages/jax/linear_util.py](https://localhost:8080/#) in call_wrapped(self, *args, **kwargs)
    164 
    165     try:
--> 166       ans = self.f(*args, **dict(self.params, **kwargs))
    167     except:
    168       # Some transformations yield from inside context managers, so we have to

[/usr/local/lib/python3.7/dist-packages/transformers/models/bart/modeling_flax_bart.py](https://localhost:8080/#) in init_weights(self, rng, input_shape)
    937             rngs,
    938             input_ids,
--> 939             attention_mask,
    940             decoder_input_ids,
    941             decoder_attention_mask,

AttributeError: module 'jax.ops' has no attribute 'index_update'

You should install transfomers master from source since this this fix is not included in the release yet. It will be available in the next release. Install transformers using
`pip install git+https://github.com/huggingface/transformers.git`

Ok perfect, that fixed it. Thank you!