# 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]:


import dataclasses
import jax
import jax.numpy as jnp
import numpy as np
from typing import Union
from alphafold.model.geometry import struct_of_array
from alphafold.model.geometry import utils
    """Compute dot product between 'self' and 'other'."""
    return self.x * other.x + self.y * other.y + self.z * other.z

  def norm(self, epsilon: float = 1e-6) -> Float:
    """Compute Norm of Vec3Array, clipped to epsilon."""
    # To avoid NaN on the backward pass, we must use maximum before the sqrt
    norm2 = self.dot(self)
    if epsilon:
      norm2 = jnp.maximum(norm2, epsilon**2)
    return jnp.sqrt(norm2)

  def norm2(self):
    return self.dot(self)

  def normalized(self, epsilon: float = 1e-6) -> Vec3Array:
    """Return unit vector with optional clipping."""
    return self / self.norm(epsilon)

  @classmethod
  def zeros(cls, shape, dtype=jnp.float32):
    """Return Vec3Array corresponding to zeros of given shape."""
    return cls(
        jnp.zeros(shape, dtype), jnp.zeros(shape, dtype),
        jnp.zeros(shape, dtype))

  def to_array(self) -> jnp.ndarray:
    return jnp.stack([self.x, self.y, self.z], axis=-1)

  @classmethod
  def from_array(cls, array):

# 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
import dataclasses
import jax
import jax.numpy as jnp
import numpy as np
from typing import Union
from alphafold.model.geometry import struct_of_array
from alphafold.model.geometry import utils

    """Compute dot product between 'self' and 'other'."""
    return self.x * other.x + self.y * other.y + self.z * other.z

  def norm(self, epsilon: float = 1e-6) -> Float:
    """Compute Norm of Vec3Array, clipped to epsilon."""
    # To avoid NaN on the backward pass, we must use maximum before the sqrt
    norm2 = self.dot(self)
    if epsilon:
      norm2 = jnp.maximum(norm2, epsilon**2)
    return jnp.sqrt(norm2)

  def norm2(self):
    return self.dot(self)

  def normalized(self, epsilon: float = 1e-6) -> Vec3Array:
    """Return unit vector with optional clipping."""
    return self / self.norm(epsilon)

  @classmethod
  def zeros(cls, shape, dtype=jnp.float32):
    """Return Vec3Array corresponding to zeros of given shape."""
    return cls(
        jnp.zeros(shape, dtype), jnp.zeros(shape, dtype),
        jnp.zeros(shape, dtype))

  def to_array(self) -> jnp.ndarray:
    return jnp.stack([self.x, self.y, self.z], axis=-1)

  @classmethod
  def from_array(cls, array):

    return cls(*utils.unstack(array))