Skip to content

cagpjax.operators.diag_like

Functions:

  • diag_like –

    Create a diagonal operator with the same shape, dtype, and device as the operator.

diag_like 

diag_like(operator: LinearOperator, values: ScalarFloat | Float[Array, N]) -> Diagonal | ScalarMul

Create a diagonal operator with the same shape, dtype, and device as the operator.

Parameters:

  • operator

    (LinearOperator) –

    Linear operator.

  • values

    (ScalarFloat | Float[Array, N]) –

    Scalar for a scalar matrix or array of diagonal values for a diagonal matrix.

Returns:

Source code in src/cagpjax/operators/diag_like.py 
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
def diag_like(
    operator: LinearOperator, values: ScalarFloat | Float[Array, "N"]
) -> Diagonal | ScalarMul:
    """Create a diagonal operator with the same shape, dtype, and device as the operator.

    Args:
        operator: Linear operator.
        values: Scalar for a scalar matrix or array of diagonal values for a diagonal matrix.

    Returns:
            Diagonal or scalar operator.
    """
    device = operator.device
    dtype = operator.dtype
    if jnp.isscalar(values):
        return ScalarMul(values, operator.shape, dtype=dtype, device=device)
    else:
        values = values.astype(dtype).to_device(device)
        return Diagonal(values)