neuralqx.operators.computational.Euclidean4d.jax.vertex_TRC module

class EuclideanVertexTRCOperatorJax(lqx, vertex, *, apply_lapse=True)

Bases: ComputationalJaxOperator

JAX-native version of EuclideanVertexTRCOperatorFast that:
  • is a PyTree (works inside jit/vmap/pmap)

  • uses a top-level jitted kernel (method itself is not jitted)

  • keeps all arrays in JAX

property max_conn_size: int

The maximum number of non zero ⟨x|O|x’⟩ for every x.

property is_hermitian: bool

Returns true if this operator is hermitian.

property dtype

The dtype of the operator’s matrix elements ⟨σ|Ô|σ’⟩.

tree_flatten()
classmethod tree_unflatten(aux, leaves)
class EuclideanVertexTRCAdjointOperatorJax(lqx, vertex, *, apply_lapse=True)

Bases: EuclideanVertexTRCOperatorJax

JAX-native adjoint of EuclideanVertexTRCOperatorJax. Only the top-level kernel changes to evaluate ΔV on branch states and flip signs.

ThiemannRegularisedVertexConstraintJax(lqx, vertex, *, apply_lapse=True, adjoint=False)