def compute_jvp_of_vector_valued_fn_wrt_activations(
    vector_valued_fn: Callable[[torch.Tensor], torch.Tensor],
    resid: torch.Tensor,
    write_vector: torch.Tensor,
) -> torch.Tensor:
    """
    vector_valued_fn takes a vector input and returns a vector. This function evaluates the jacobian
    of the vector_valued_fn with respect to the vector input, at the vector specified by resid, and
    then multiplies the jacobian by write_vector. pytorch's jvp function is used to perform this
    efficiently (without instantiating the full jacobian matrix). This can be considered a "forward
    prop" or the multiplication of a vector by the derivative of a computation graph in the forward
    direction (rather than the more common reverse direction)
    """
    jacobian_vector_product = torch_jvp(
        vector_valued_fn,
        (resid,),
        (write_vector,),
    )[1]

    return jacobian_vector_product

    vector_valued_fn takes a vector input and returns a vector. This function evaluates the jacobian
    of the vector_valued_fn with respect to the vector input, at the vector specified by resid, and
    then multiplies the jacobian by write_vector. pytorch's jvp function is used to perform this
    efficiently (without instantiating the full jacobian matrix). This can be considered a "forward
    prop" or the multiplication of a vector by the derivative of a computation graph in the forward
    direction (rather than the more common reverse direction)
    
neuron_explainer/activations/derived_scalars/reconstituter_class.py