Repository: Genesis-Embodied-AI/Genesis
File: genesis/utils/array_class.py (Python)

Code before:
    cg_beta: V_ANNOTATION
    cg_pg_dot_pMg: V_ANNOTATION
    # Optional Newton fields
    nt_H: V_ANNOTATION
    nt_vec: V_ANNOTATION


def get_constraint_state(constraint_solver, solver):
    f_batch = solver._batch_shape
    len_constraints = constraint_solver.len_constraints
    len_constraints_ = constraint_solver.len_constraints_

    jac_shape = solver._batch_shape((len_constraints_, solver.n_dofs_))
    if (jac_shape[0] * jac_shape[1] * jac_shape[2]) > np.iinfo(np.int32).max:
        raise ValueError(
            f"Jacobian shape {jac_shape} is too large for int32. "
            "Consider reducing the number of constraints or the number of degrees of freedom."
        )

    if solver._static_rigid_sim_config.noslip_iterations > 0:
        if len_constraints_ * len_constraints_ * f_batch()[0] > 2 * 10**9:
            gs.logger.warning(
                f"efc_AR shape {len_constraints_}x{len_constraints_}x{f_batch()[0]} is very large. "
                f"Consider manually set a smaller 'max_collision_pairs' in RigidOptions to reduce the size of "
                "reserved memory. "
            )
        efc_AR_shape = solver._batch_shape((len_constraints_, len_constraints_))
        efc_b_shape = solver._batch_shape(len_constraints_)
    else:
        efc_AR_shape = 1
        efc_b_shape = 1

    kwargs = {
        "n_constraints": V(dtype=gs.ti_int, shape=f_batch()),
        "ti_n_equalities": V(gs.ti_int, shape=solver._batch_shape()),
        "jac": V(dtype=gs.ti_float, shape=solver._batch_shape((len_constraints_, solver.n_dofs_))),
        "diag": V(dtype=gs.ti_float, shape=solver._batch_shape(len_constraints_)),
        "aref": V(dtype=gs.ti_float, shape=solver._batch_shape(len_constraints_)),
        "jac_relevant_dofs": V(gs.ti_int, shape=solver._batch_shape((len_constraints_, solver.n_dofs_))),
        "jac_n_relevant_dofs": V(gs.ti_int, shape=solver._batch_shape(len_constraints_)),
        "n_constraints_equality": V(gs.ti_int, shape=solver._batch_shape()),
        "n_constraints_frictionloss": V(gs.ti_int, shape=solver._batch_shape()),
        "improved": V(gs.ti_int, shape=solver._batch_shape()),
        "Jaref": V(dtype=gs.ti_float, shape=solver._batch_shape(len_constraints_)),
        "Ma": V(dtype=gs.ti_float, shape=solver._batch_shape(solver.n_dofs_)),
        "Ma_ws": V(dtype=gs.ti_float, shape=solver._batch_shape(solver.n_dofs_)),
        "grad": V(dtype=gs.ti_float, shape=solver._batch_shape(solver.n_dofs_)),
        "Mgrad": V(dtype=gs.ti_float, shape=solver._batch_shape(solver.n_dofs_)),
        "search": V(dtype=gs.ti_float, shape=solver._batch_shape(solver.n_dofs_)),
        "efc_D": V(dtype=gs.ti_float, shape=solver._batch_shape(len_constraints_)),
        "efc_frictionloss": V(dtype=gs.ti_float, shape=solver._batch_shape(len_constraints_)),

Review comment (performance):
Only move to new line if necessary:
```python
            gs.logger.warning(
                f"efc_AR shape {len_constraints_}x{len_constraints_}x{f_batch()[0]} is very large. Consider manually "
                f"set a smaller 'max_collision_pairs' in RigidOptions to reduce the size of reserved memory. "
            )
```

Code after:
    cg_beta: V_ANNOTATION
    cg_pg_dot_pMg: V_ANNOTATION
    # Optional Newton fields
    nt_H: V_ANNOTATION
    nt_vec: V_ANNOTATION


def get_constraint_state(constraint_solver, solver):
    f_batch = solver._batch_shape
    len_constraints = constraint_solver.len_constraints
    len_constraints_ = constraint_solver.len_constraints_

    jac_shape = solver._batch_shape((len_constraints_, solver.n_dofs_))
    if (jac_shape[0] * jac_shape[1] * jac_shape[2]) > np.iinfo(np.int32).max:
        raise ValueError(
            f"Jacobian shape {jac_shape} is too large for int32. "
            "Consider reducing the number of constraints or the number of degrees of freedom."
        )

    if solver._static_rigid_sim_config.noslip_iterations > 0:
        if len_constraints_ * len_constraints_ * f_batch()[0] > 2 * 10**9:
            gs.logger.warning(
                f"efc_AR shape {len_constraints_}x{len_constraints_}x{f_batch()[0]} is very large. Consider manually "
                f"set a smaller 'max_collision_pairs' in RigidOptions to reduce the size of reserved memory. "
            )
        efc_AR_shape = solver._batch_shape((len_constraints_, len_constraints_))
        efc_b_shape = solver._batch_shape(len_constraints_)
    else:
        efc_AR_shape = 1
        efc_b_shape = 1

    kwargs = {
        "n_constraints": V(dtype=gs.ti_int, shape=f_batch()),
        "ti_n_equalities": V(gs.ti_int, shape=solver._batch_shape()),
        "jac": V(dtype=gs.ti_float, shape=solver._batch_shape((len_constraints_, solver.n_dofs_))),
        "diag": V(dtype=gs.ti_float, shape=solver._batch_shape(len_constraints_)),
        "aref": V(dtype=gs.ti_float, shape=solver._batch_shape(len_constraints_)),
        "jac_relevant_dofs": V(gs.ti_int, shape=solver._batch_shape((len_constraints_, solver.n_dofs_))),
        "jac_n_relevant_dofs": V(gs.ti_int, shape=solver._batch_shape(len_constraints_)),
        "n_constraints_equality": V(gs.ti_int, shape=solver._batch_shape()),
        "n_constraints_frictionloss": V(gs.ti_int, shape=solver._batch_shape()),
        "improved": V(gs.ti_int, shape=solver._batch_shape()),
        "Jaref": V(dtype=gs.ti_float, shape=solver._batch_shape(len_constraints_)),
        "Ma": V(dtype=gs.ti_float, shape=solver._batch_shape(solver.n_dofs_)),
        "Ma_ws": V(dtype=gs.ti_float, shape=solver._batch_shape(solver.n_dofs_)),
        "grad": V(dtype=gs.ti_float, shape=solver._batch_shape(solver.n_dofs_)),
        "Mgrad": V(dtype=gs.ti_float, shape=solver._batch_shape(solver.n_dofs_)),
        "search": V(dtype=gs.ti_float, shape=solver._batch_shape(solver.n_dofs_)),
        "efc_D": V(dtype=gs.ti_float, shape=solver._batch_shape(len_constraints_)),
        "efc_frictionloss": V(dtype=gs.ti_float, shape=solver._batch_shape(len_constraints_)),
        "efc_force": V(dtype=gs.ti_float, shape=solver._batch_shape(len_constraints_)),