class_docstr: Describes structure present in a primitive derivative dy/dw.
  
  # TODO(romann): make this a python dataclass.
  
  Attributes:
    out_trace:
      axes of the primitive `y` output along which the primitive Jacobian
      `dy/dw` is constant-block diagonal along the respective axes in the input
      `in_trace`.
  
    in_trace:
      axes of the primitive `y` inputs along which the primitive Jacobian
      `dy/dw` is constant-block diagonal along the respective axes in the output
      `out_trace`.
  
    in_trace_idxs:
      indices of input variables to which `in_trace` axes are applied. Other
      variables are considered untouched.
  
    out_diagonal:
      axes of the primitive `y` output along which the primitive Jacobian
      `dy/dw` is (not constant) block diagonal along the respective axes in the
      input `in_diagonal`.
  
    in_diagonal:
      axes of the primitive `y` inputs along which the primitive Jacobian
      `dy/dw` is (not constant) block diagonal along the respective axes in the
      output `out_diagonal`. Each entry in the `in_diagonal` tuple is a tuple of
      length equal to the number of input variables; each entry in the tuple is
      either an integer axis number correspomnding to the respective input
      variable, or `None`, meaning that the respective variable is considered
      untouched.
  
    out_broadcast:
      axes of the primitive `y` output along which the primitive Jacobian
      `dy/dw` is block-tiled.
  
    out_broadcast_idxs:
      indices of input variables that need to be squeezed along the
      `out_broadcast` axes in order for the primitive `y` to return the slice
      that is being tiled along `out_broadcast` in the full output.
  
    in_broadcast:
      axes of the primitive `y` inputs along which the primitive Jacobian
      `dy/dw` is block-tiled.
  
    in_broadcast_idx:
      indices of input variables that need to be squeezed along the
      `in_broadcast` axes in order for the primitive Jacobian `dy/dw` to return
      the slice that is being tiled along `in_broadcast` in the full output.
docstr: Defines interaction with structure of the other primitive dy2/dw.
docstr_tok: ["Defines","interaction","with","structure","of","the","other","primitive","dy2\/dw","."]
code: def __and__(self, other):
      """Defines interaction with structure of the other primitive dy2/dw."""
      assert len(self.in_trace) == len(self.out_trace), (self, other)
      assert len(other.in_trace) == len(other.out_trace), (self, other)
  
      in_trace_idxs = self.in_trace_idxs
      in_trace = tuple(i for i in self.in_trace if i in other.in_trace)
  
      out_trace = tuple(
          self.out_trace[i]
          for i in range(len(self.out_trace))
          if self.in_trace[i] in other.in_trace
      )
  
      assert len(in_trace) == len(out_trace), (self, other)
  
      out_diagonal = tuple(
          i for i in self.out_diagonal if i in other.out_diagonal
      )
      in_diagonal = tuple(
          i
          for ix, i in enumerate(self.in_diagonal)
          if self.out_diagonal[ix] in other.out_diagonal
      )
  
      out_broadcast = tuple(
          i for i in self.out_broadcast if i in other.out_broadcast
      )
  
      in_broadcast = tuple(
          i for i in self.out_broadcast if i in other.out_broadcast
      )
  
      return Structure(
          out_trace=out_trace,
          in_trace=in_trace,
          in_trace_idxs=in_trace_idxs,
          out_diagonal=out_diagonal,
          in_diagonal=in_diagonal,
          out_broadcast=out_broadcast,
          out_broadcast_idxs=self.out_broadcast_idxs,
          in_broadcast=in_broadcast,
          in_broadcast_idx=self.in_broadcast_idx,
      )
  
imports: import functools
  from typing import Any, Callable, Optional, Union
  import jax
  from jax import lax
  from jax.core import AbstractValue
  from jax.core import Jaxpr
  from jax.core import JaxprEqn
  from jax.core import Literal
  from jax.core import Primitive
  from jax.core import ShapedArray
  from jax.core import Var
  from jax.interpreters import ad
  import jax.numpy
  import numpy
  from .None import utils
  from .dataclasses import dataclass
  from .dataclasses import field
label_desc: Use image node_id 1 for calling the Structure obj's underlying member method code with example usage: obj.__and__(other) and returns: Structure