class_docstr: Protein folding specific Linear Module.
  
  This differs from the standard Haiku Linear in a few ways:
      * It supports inputs of arbitrary rank
      * Initializers are specified by strings
  
  This code is adapted from DeepMind's AlphaFold code release
  (https://github.com/deepmind/alphafold).
  
  Examples
  --------
  >>> import deepchem as dc
  >>> import haiku as hk
  >>> import jax
  >>> import deepchem.models.jax_models.layers
  >>> def forward_model(x):
  ...   layer = dc.models.jax_models.layers.Linear(2)
  ...   return layer(x)
  >>> f = hk.transform(forward_model)
  >>> rng = jax.random.PRNGKey(42)
  >>> x = jnp.ones([8, 28 * 28])
  >>> params = f.init(rng, x)
  >>> output = f.apply(params, rng, x)
class_docstr_tok: ["Protein","folding","specific","Linear","Module",".","This","differs","from","the","standard","Haiku","Linear","in","a","few","ways",":","*","It","supports","inputs","of","arbitrary","rank","*","Initializers","are","specified","by","strings","This","code","is","adapted","from","DeepMind","'s","AlphaFold","code","release","(","https",":","\/\/github.com\/deepmind\/alphafold",")",".","Examples","--","--","--","--",">",">",">","import","deepchem","as","dc",">",">",">","import","haiku","as","hk",">",">",">","import","jax",">",">",">","import","deepchem.models.jax_models.layers",">",">",">","def","forward_model","(","x",")",":","...","layer","=","dc.models.jax_models.layers.Linear","(","2",")","...","return","layer","(","x",")",">",">",">","f","=","hk.transform","(","forward_model",")",">",">",">","rng","=","jax.random.PRNGKey","(","42",")",">",">",">","x","=","jnp.ones","(","[","8",",","28","*","28","]",")",">",">",">","params","=","f.init","(","rng",",","x",")",">",">",">","output","=","f.apply","(","params",",","rng",",","x",")"]
docstr: Connects Module.
  
  Parameters
  ----------
  inputs: jnp.ndarray
      Tensor of shape [..., num_channel]
  
  Returns
  -------
  output of shape [..., num_output]
docstr_tok: ["Connects","Module",".","Parameters","--","--","--","--","--","inputs",":","jnp.ndarray","Tensor","of","shape","[","...",",","num_channel","]","Returns","--","--","--","-","output","of","shape","[","...",",","num_output","]"]
code: def __call__(self, inputs: jnp.ndarray) -> jnp.ndarray:
      """Connects Module.
  
      Parameters
      ----------
      inputs: jnp.ndarray
          Tensor of shape [..., num_channel]
  
      Returns
      -------
      output of shape [..., num_output]
      """
      n_channels = int(inputs.shape[-1])
  
      weight_shape = [n_channels, self.num_output]
      if self.initializer == "linear":
          weight_init = hk.initializers.VarianceScaling(
              mode="fan_in", scale=1.0
          )
      elif self.initializer == "relu":
          weight_init = hk.initializers.VarianceScaling(
              mode="fan_in", scale=2.0
          )
      elif self.initializer == "zeros":
          weight_init = hk.initializers.Constant(0.0)
  
      weights = hk.get_parameter(
          "weights", weight_shape, inputs.dtype, weight_init
      )
  
      # this is equivalent to einsum('...c,cd->...d', inputs, weights)
      # but turns out to be slightly faster
      inputs = jnp.swapaxes(inputs, -1, -2)
      output = jnp.einsum("...cb,cd->...db", inputs, weights)
      output = jnp.swapaxes(output, -1, -2)
  
      if self.use_bias:
          bias = hk.get_parameter(
              "bias",
              [self.num_output],
              inputs.dtype,
              hk.initializers.Constant(self.bias_init),
          )
          output += bias
  
      return output
  
imports: import logging
label_desc: Use image node_id 2 for calling the Linear obj's underlying member method code with example usage: obj.__call__(inputs) and returns: output