class_docstr: Token write operations motivated by the `write' in Neural Turing Machines.
  
  Instead of directly using the token summarization (with TokenLearner), it uses
  a similar but different mechanism to (soft-)select memory elements to zero out
  and write to them. This can be used as an alternative write operation in the
  TTM, particularly when the memory size is huge.
code: def __call__(
      self,
      memory: jnp.ndarray,
      control_inputs: jnp.ndarray,
      training: bool = False,
  ) -> jnp.ndarray:
      selected = nn.LayerNorm()(memory)
  
      selected = attention_layers.MlpBlock(
          mlp_dim=self.bottleneck_dim,
          out_dim=self.num_tokens,
          dropout_rate=self.dropout_rate,
          activation_fn=nn.gelu,
      )(selected, deterministic=not training)
  
      selected = jnp.transpose(
          selected, [0, 2, 1]
      )  # Shape: [bs, n_token, hw].
      selected = jax.nn.softmax(selected, axis=-1)
  
      et = nn.LayerNorm()(control_inputs)
      et = jnp.transpose(et, [0, 2, 1])  # Shape: [bs, c, hw].
      et = attention_layers.MlpBlock(
          mlp_dim=self.bottleneck_dim,
          out_dim=self.num_tokens,
          dropout_rate=self.dropout_rate,
          activation_fn=nn.gelu,
      )(
          et, deterministic=not training
      )  # Shape: [bs, c, n_token].
      et = jnp.transpose(et, [0, 2, 1])  # Shape: [bs, n_token, c].
  
      et = attention_layers.MlpBlock(
          mlp_dim=self.bottleneck_dim,
          dropout_rate=self.dropout_rate,
          activation_fn=nn.gelu,
      )(et, deterministic=not training)
  
      wet = jnp.expand_dims(selected, -1) * jnp.expand_dims(
          et, 2
      )  # Shape: [bs, n_token, hw, c].
      wet = 1 - wet
      wet = jnp.prod(wet, axis=1)  # Shape: [bs, hw, c].
  
      output = memory * wet
  
      at = nn.LayerNorm()(control_inputs)
      at = jnp.transpose(at, [0, 2, 1])  # Shape: [bs, c, hw].
      at = attention_layers.MlpBlock(
          mlp_dim=self.bottleneck_dim,
          out_dim=self.num_tokens,
          dropout_rate=self.dropout_rate,
          activation_fn=nn.gelu,
      )(
          at, deterministic=not training
      )  # Shape: [bs, c, n_token].
      at = jnp.transpose(at, [0, 2, 1])  # Shape: [bs, n_token, c].
  
      at = attention_layers.MlpBlock(
          mlp_dim=self.bottleneck_dim,
          dropout_rate=self.dropout_rate,
          activation_fn=nn.gelu,
      )(at, deterministic=not training)
  
      wat = jnp.expand_dims(selected, -1) * jnp.expand_dims(
          at, 2
      )  # Shape: [bs, n_token, hw, c].
      wat = 1 - wat
      wat = jnp.mean(wat, axis=1)  # Shape: [bs, hw, c].
  
      output += wat
  
      return output
  
imports: from typing import Tuple
  from absl import logging
  import flax.linen
  import jax
  import jax.numpy
  from scenic.model_lib.layers import attention_layers
  from scenic.projects.baselines import mixer
  from scenic.projects.baselines import vit
  from scenic.projects.token_learner import model
label_desc: Use image node_id 1 for calling the TokenAddEraseWrite obj's underlying member method code with example usage: obj.__call__(memory, control_inputs, training) and returns: output