# Path: alphafold/model/geometry/struct_of_array.py
# def get_item(instance, key):
# def get_shape(instance):
# def get_len(instance):
# def get_dtype(instance):
# def replace(instance, **kwargs):
# def post_init(instance):
# def flatten(instance):
# def make_metadata_class(cls):
# def get_fields(cls_or_instance, filterfn, return_values=False):
# def get_array_fields(cls, return_values=False):
# def get_metadata_fields(cls, return_values=False):
#   def __init__(self, same_dtype=True):
#   def __call__(self, cls):
#     def unflatten(aux, data):
# class StructOfArray:
# 
# Path: alphafold/model/geometry/utils.py
# def unstack(value: jnp.ndarray, axis: int = -1) -> List[jnp.ndarray]:
# 
# Path: alphafold/model/geometry/vector.py
# VERSION = '0.1'
# class Vec3Array:
#   def __post_init__(self):
#   def __add__(self, other: Vec3Array) -> Vec3Array:
#   def __sub__(self, other: Vec3Array) -> Vec3Array:
#   def __mul__(self, other: Float) -> Vec3Array:
#   def __rmul__(self, other: Float) -> Vec3Array:
#   def __truediv__(self, other: Float) -> Vec3Array:
#   def __neg__(self) -> Vec3Array:
#   def __pos__(self) -> Vec3Array:
#   def cross(self, other: Vec3Array) -> Vec3Array:
#   def dot(self, other: Vec3Array) -> Float:
#   def norm(self, epsilon: float = 1e-6) -> Float:
#   def norm2(self):
#   def normalized(self, epsilon: float = 1e-6) -> Vec3Array:
#   def zeros(cls, shape, dtype=jnp.float32):
#   def to_array(self) -> jnp.ndarray:
#   def from_array(cls, array):
#   def __getstate__(self):
#   def __setstate__(self, state):
# def square_euclidean_distance(vec1: Vec3Array,
#                               vec2: Vec3Array,
#                               epsilon: float = 1e-6) -> Float:
# def dot(vector1: Vec3Array, vector2: Vec3Array) -> Float:
# def cross(vector1: Vec3Array, vector2: Vec3Array) -> Float:
# def norm(vector: Vec3Array, epsilon: float = 1e-6) -> Float:
# def normalized(vector: Vec3Array, epsilon: float = 1e-6) -> Vec3Array:
# def euclidean_distance(vec1: Vec3Array,
#                        vec2: Vec3Array,
#                        epsilon: float = 1e-6) -> Float:
# def dihedral_angle(a: Vec3Array, b: Vec3Array, c: Vec3Array,
#                    d: Vec3Array) -> Float:
# def random_gaussian_vector(shape, key, dtype=jnp.float32):


import dataclasses
import jax
import jax.numpy as jnp
import numpy as np
from alphafold.model.geometry import struct_of_array
from alphafold.model.geometry import utils
from alphafold.model.geometry import vector
    ones = jnp.ones(shape, dtype=dtype)
    zeros = jnp.zeros(shape, dtype=dtype)
    return cls(ones, zeros, zeros, zeros, ones, zeros, zeros, zeros, ones)

  @classmethod
  def from_two_vectors(cls, e0: vector.Vec3Array,
                       e1: vector.Vec3Array) -> Rot3Array:
    """Construct Rot3Array from two Vectors.

    Rot3Array is constructed such that in the corresponding frame 'e0' lies on
    the positive x-Axis and 'e1' lies in the xy plane with positive sign of y.

    Args:
      e0: Vector
      e1: Vector
    Returns:
      Rot3Array
    """
    # Normalize the unit vector for the x-axis, e0.
    e0 = e0.normalized()
    # make e1 perpendicular to e0.
    c = e1.dot(e0)
    e1 = (e1 - c * e0).normalized()
    # Compute e2 as cross product of e0 and e1.
    e2 = e0.cross(e1)
    return cls(e0.x, e1.x, e2.x, e0.y, e1.y, e2.y, e0.z, e1.z, e2.z)

  @classmethod
  def from_array(cls, array: jnp.ndarray) -> Rot3Array:
    """Construct Rot3Array Matrix from array of shape. [..., 3, 3]."""

# Path: alphafold/model/geometry/struct_of_array.py
# def get_item(instance, key):
# def get_shape(instance):
# def get_len(instance):
# def get_dtype(instance):
# def replace(instance, **kwargs):
# def post_init(instance):
# def flatten(instance):
# def make_metadata_class(cls):
# def get_fields(cls_or_instance, filterfn, return_values=False):
# def get_array_fields(cls, return_values=False):
# def get_metadata_fields(cls, return_values=False):
#   def __init__(self, same_dtype=True):
#   def __call__(self, cls):
#     def unflatten(aux, data):
# class StructOfArray:
# 
# Path: alphafold/model/geometry/utils.py
# def unstack(value: jnp.ndarray, axis: int = -1) -> List[jnp.ndarray]:
# 
# Path: alphafold/model/geometry/vector.py
# VERSION = '0.1'
# class Vec3Array:
#   def __post_init__(self):
#   def __add__(self, other: Vec3Array) -> Vec3Array:
#   def __sub__(self, other: Vec3Array) -> Vec3Array:
#   def __mul__(self, other: Float) -> Vec3Array:
#   def __rmul__(self, other: Float) -> Vec3Array:
#   def __truediv__(self, other: Float) -> Vec3Array:
#   def __neg__(self) -> Vec3Array:
#   def __pos__(self) -> Vec3Array:
#   def cross(self, other: Vec3Array) -> Vec3Array:
#   def dot(self, other: Vec3Array) -> Float:
#   def norm(self, epsilon: float = 1e-6) -> Float:
#   def norm2(self):
#   def normalized(self, epsilon: float = 1e-6) -> Vec3Array:
#   def zeros(cls, shape, dtype=jnp.float32):
#   def to_array(self) -> jnp.ndarray:
#   def from_array(cls, array):
#   def __getstate__(self):
#   def __setstate__(self, state):
# def square_euclidean_distance(vec1: Vec3Array,
#                               vec2: Vec3Array,
#                               epsilon: float = 1e-6) -> Float:
# def dot(vector1: Vec3Array, vector2: Vec3Array) -> Float:
# def cross(vector1: Vec3Array, vector2: Vec3Array) -> Float:
# def norm(vector: Vec3Array, epsilon: float = 1e-6) -> Float:
# def normalized(vector: Vec3Array, epsilon: float = 1e-6) -> Vec3Array:
# def euclidean_distance(vec1: Vec3Array,
#                        vec2: Vec3Array,
#                        epsilon: float = 1e-6) -> Float:
# def dihedral_angle(a: Vec3Array, b: Vec3Array, c: Vec3Array,
#                    d: Vec3Array) -> Float:
# def random_gaussian_vector(shape, key, dtype=jnp.float32):


# Path: alphafold/model/geometry/rotation_matrix.py
import dataclasses
import jax
import jax.numpy as jnp
import numpy as np
from alphafold.model.geometry import struct_of_array
from alphafold.model.geometry import utils
from alphafold.model.geometry import vector

    ones = jnp.ones(shape, dtype=dtype)
    zeros = jnp.zeros(shape, dtype=dtype)
    return cls(ones, zeros, zeros, zeros, ones, zeros, zeros, zeros, ones)

  @classmethod
  def from_two_vectors(cls, e0: vector.Vec3Array,
                       e1: vector.Vec3Array) -> Rot3Array:
    """Construct Rot3Array from two Vectors.

    Rot3Array is constructed such that in the corresponding frame 'e0' lies on
    the positive x-Axis and 'e1' lies in the xy plane with positive sign of y.

    Args:
      e0: Vector
      e1: Vector
    Returns:
      Rot3Array
    """
    # Normalize the unit vector for the x-axis, e0.
    e0 = e0.normalized()
    # make e1 perpendicular to e0.
    c = e1.dot(e0)
    e1 = (e1 - c * e0).normalized()
    # Compute e2 as cross product of e0 and e1.
    e2 = e0.cross(e1)
    return cls(e0.x, e1.x, e2.x, e0.y, e1.y, e2.y, e0.z, e1.z, e2.z)

  @classmethod
  def from_array(cls, array: jnp.ndarray) -> Rot3Array:
    """Construct Rot3Array Matrix from array of shape. [..., 3, 3]."""

    unstacked = utils.unstack(array, axis=-2)