class_docstr: A simple MLP.
  
  Attributes:
    nerf_trunk_depth: int, the depth of the first part of MLP.
    nerf_trunk_width: int, the width of the first part of MLP.
    nerf_rgb_branch_depth: int, the depth of the second part of MLP.
    nerf_rgb_branch_width: int, the width of the second part of MLP.
    activation: function, the activation function used in the MLP.
    skips: which layers to add skip layers to.
    alpha_channels: int, the number of alpha_channelss.
    rgb_channels: int, the number of rgb_channelss.
    condition_density: if True put the condition at the begining which
      conditions the density of the field.
class_docstr_tok: ["A","simple","MLP",".","Attributes",":","nerf_trunk_depth",":","int",",","the","depth","of","the","first","part","of","MLP",".","nerf_trunk_width",":","int",",","the","width","of","the","first","part","of","MLP",".","nerf_rgb_branch_depth",":","int",",","the","depth","of","the","second","part","of","MLP",".","nerf_rgb_branch_width",":","int",",","the","width","of","the","second","part","of","MLP",".","activation",":","function",",","the","activation","function","used","in","the","MLP",".","skips",":","which","layers","to","add","skip","layers","to",".","alpha_channels",":","int",",","the","number","of","alpha_channelss",".","rgb_channels",":","int",",","the","number","of","rgb_channelss",".","condition_density",":","if","True","put","the","condition","at","the","begining","which","conditions","the","density","of","the","field","."]
docstr: Multi-layer perception for nerf.
  
  Args:
    x: sample points with shape [batch, num_coarse_samples, feature].
    trunk_condition: a condition array provided to the trunk.
    alpha_condition: a condition array provided to the alpha branch.
    rgb_condition: a condition array provided in the RGB branch.
  
  Returns:
    raw: [batch, num_coarse_samples, rgb_channels+alpha_channels].
docstr_tok: ["Multi-layer","perception","for","nerf",".","Args",":","x",":","sample","points","with","shape","[","batch",",","num_coarse_samples",",","feature","]",".","trunk_condition",":","a","condition","array","provided","to","the","trunk",".","alpha_condition",":","a","condition","array","provided","to","the","alpha","branch",".","rgb_condition",":","a","condition","array","provided","in","the","RGB","branch",".","Returns",":","raw",":","[","batch",",","num_coarse_samples",",","rgb_channels+alpha_channels","]","."]
code: def __call__(
      self, x, trunk_condition, alpha_condition, rgb_condition
  ):
      """Multi-layer perception for nerf.
  
      Args:
        x: sample points with shape [batch, num_coarse_samples, feature].
        trunk_condition: a condition array provided to the trunk.
        alpha_condition: a condition array provided to the alpha branch.
        rgb_condition: a condition array provided in the RGB branch.
  
      Returns:
        raw: [batch, num_coarse_samples, rgb_channels+alpha_channels].
      """
      dense = functools.partial(
          nn.Dense, kernel_init=jax.nn.initializers.glorot_uniform()
      )
  
      feature_dim = x.shape[-1]
      num_samples = x.shape[1]
      x = x.reshape([-1, feature_dim])
  
      def broadcast_condition(c):
          # Broadcast condition from [batch, feature] to
          # [batch, num_coarse_samples, feature] since all the samples along the
          # same ray has the same viewdir.
          c = jnp.tile(c[:, None, :], (1, num_samples, 1))
          # Collapse the [batch, num_coarse_samples, feature] tensor to
          # [batch * num_coarse_samples, feature] to be fed into nn.Dense.
          c = c.reshape([-1, c.shape[-1]])
          return c
  
      trunk_mlp = MLP(
          depth=self.trunk_depth,
          width=self.trunk_width,
          hidden_activation=self.activation,
          hidden_init=jax.nn.initializers.glorot_uniform(),
          skips=self.skips,
      )
      rgb_mlp = MLP(
          depth=self.rgb_branch_depth,
          width=self.rgb_branch_width,
          hidden_activation=self.activation,
          hidden_init=jax.nn.initializers.glorot_uniform(),
          output_init=jax.nn.initializers.glorot_uniform(),
          output_channels=self.rgb_channels,
      )
      alpha_mlp = MLP(
          depth=self.alpha_branch_depth,
          width=self.alpha_branch_width,
          hidden_activation=self.activation,
          hidden_init=jax.nn.initializers.glorot_uniform(),
          output_init=jax.nn.initializers.glorot_uniform(),
          output_channels=self.alpha_channels,
      )
  
      if trunk_condition is not None:
          trunk_condition = broadcast_condition(trunk_condition)
          trunk_input = jnp.concatenate([x, trunk_condition], axis=-1)
      else:
          trunk_input = x
      x = trunk_mlp(trunk_input)
  
      if (alpha_condition is not None) or (rgb_condition is not None):
          bottleneck = dense(self.trunk_width, name="bottleneck")(x)
  
      if alpha_condition is not None:
          alpha_condition = broadcast_condition(alpha_condition)
          alpha_input = jnp.concatenate(
              [bottleneck, alpha_condition], axis=-1
          )
      else:
          alpha_input = x
      alpha = alpha_mlp(alpha_input)
  
      if rgb_condition is not None:
          rgb_condition = broadcast_condition(rgb_condition)
          rgb_input = jnp.concatenate(
              [bottleneck, rgb_condition], axis=-1
          )
      else:
          rgb_input = x
      rgb = rgb_mlp(rgb_input)
  
      return {
          "rgb": rgb.reshape((-1, num_samples, self.rgb_channels)),
          "alpha": alpha.reshape(
              (-1, num_samples, self.alpha_channels)
          ),
      }
  
imports: import functools
  from typing import Optional, Tuple
  from flax import linen
  import jax
  import jax.numpy
  from nerfies import types
label_desc: Use image node_id 1 for calling the NerfMLP obj's underlying member method code with example usage: obj.__call__(x, trunk_condition, alpha_condition, rgb_condition) and returns: dict, c