How to use the pyhf.tensorlib.astensor function in pyhf

def test_code0_validation(backend, do_tensorized_calc):
    histogramssets = [[[[0.5], [1.0], [2.0]]]]
    alphasets = pyhf.tensorlib.astensor([[-2, -1, 0, 1, 2]])
    expected = pyhf.tensorlib.astensor([[[[0], [0.5], [1.0], [2.0], [3.0]]]])

    interpolator = pyhf.interpolators.get(0, do_tensorized_calc=do_tensorized_calc)(
        histogramssets, subscribe=False
    )
    result_deltas = pyhf.tensorlib.astensor(interpolator(alphasets))

    # calculate the actual change
    histogramssets = pyhf.tensorlib.astensor(histogramssets)
    allsets_allhistos_noms_repeated = pyhf.tensorlib.einsum(
        'sa,shb->shab',
        pyhf.tensorlib.ones(pyhf.tensorlib.shape(alphasets)),
        histogramssets[:, :, 1],
    )
    results = allsets_allhistos_noms_repeated + result_deltas

    assert (
        pytest.approx(np.asarray(pyhf.tensorlib.tolist(results)).ravel().tolist())
        == np.asarray(pyhf.tensorlib.tolist(expected)).ravel().tolist()
    )

'type': 'normfactor',
                'name': 'mu2',
                'data': {'mask': [False, True, True]},
            },
            'background': {
                'type': 'normfactor',
                'name': 'mu2',
                'data': {'mask': [False, True, True]},
            },
        },
    }
    hsc = normfactor_combined(
        [('mu1', 'normfactor'), ('mu2', 'normfactor')], mc, mega_mods
    )

    mod = hsc.apply(pyhf.tensorlib.astensor([2.0, 3.0]))
    shape = pyhf.tensorlib.shape(mod)
    assert shape == (2, 2, 1, 3)

    mod = np.asarray(pyhf.tensorlib.tolist(mod))
    assert np.allclose(mod[0, 0, 0], [2.0, 1.0, 1.0])
    assert np.allclose(mod[1, 0, 0], [1.0, 3.0, 3.0])

    hsc = normfactor_combined(
        [('mu1', 'normfactor'), ('mu2', 'normfactor')], mc, mega_mods, batch_size=4
    )

    mod = hsc.apply(
        pyhf.tensorlib.astensor([[1.0, 5.0], [2.0, 6.0], [3.0, 7.0], [4.0, 8.0]])
    )
    shape = pyhf.tensorlib.shape(mod)
    assert shape == (2, 2, 4, 3)

default_backend.normal_logpdf(data, mu, sigma)
                for data, mu, sigma in zip(
                    [0, 0, 0, 0, 0], [0, 0, 0, 0, 0], [1.5, 2.0, 1.0, 1.0, 1.0]
                )
            ]
        ),
    )
    assert result.shape == (1,)

    suggested_pars = [1.0] * 3 + [1, 2, 3, 4, 5]  # 2 pois 5 norm
    constraint = gaussian_constraint_combined(config, batch_size=1)
    result = default_backend.astensor(
        pyhf.tensorlib.tolist(
            constraint.logpdf(
                pyhf.tensorlib.astensor(config.auxdata),
                pyhf.tensorlib.astensor(suggested_pars),
            )
        )
    )
    assert np.isclose(
        result[0],
        sum(
            [
                default_backend.normal_logpdf(data, mu, sigma)
                for data, mu, sigma in zip(
                    [0, 0, 0, 0, 0], [1, 2, 3, 4, 5], [1.5, 2.0, 1.0, 1.0, 1.0]
                )
            ]
        ),
    )
    assert result.shape == (1,)

def test_code1_validation(backend, do_tensorized_calc):
    histogramssets = [[[[0.9], [1.0], [1.1]]]]
    alphasets = pyhf.tensorlib.astensor([[-2, -1, 0, 1, 2]])
    expected = pyhf.tensorlib.astensor(
        [[[[0.9 ** 2], [0.9], [1.0], [1.1], [1.1 ** 2]]]]
    )

    interpolator = pyhf.interpolators.get(1, do_tensorized_calc=do_tensorized_calc)(
        histogramssets, subscribe=False
    )
    result_deltas = interpolator(alphasets)

    # calculate the actual change
    histogramssets = pyhf.tensorlib.astensor(histogramssets)
    allsets_allhistos_noms_repeated = pyhf.tensorlib.einsum(
        'sa,shb->shab',
        pyhf.tensorlib.ones(pyhf.tensorlib.shape(alphasets)),
        histogramssets[:, :, 1],
    )
    results = allsets_allhistos_noms_repeated * result_deltas

    assert (
        pytest.approx(np.asarray(pyhf.tensorlib.tolist(results)).ravel().tolist())
        == np.asarray(pyhf.tensorlib.tolist(expected)).ravel().tolist()
    )

def test_validate_implementation(backend, interpcode, random_histosets_alphasets_pair):
    histogramssets, alphasets = random_histosets_alphasets_pair

    # single-float precision backends, calculate using single-floats
    if pyhf.tensorlib.name in ['tensorflow', 'pytorch']:
        abs_tolerance = 1e-6
        histogramssets = np.asarray(histogramssets, dtype=np.float32)
        alphasets = np.asarray(alphasets, dtype=np.float32)
    else:
        abs_tolerance = 1e-12

    histogramssets = histogramssets.tolist()
    alphasets = pyhf.tensorlib.astensor(alphasets.tolist())

    slow_interpolator = pyhf.interpolators.get(interpcode, do_tensorized_calc=False)(
        histogramssets, subscribe=False
    )
    fast_interpolator = pyhf.interpolators.get(interpcode, do_tensorized_calc=True)(
        histogramssets, subscribe=False
    )
    slow_result = np.asarray(pyhf.tensorlib.tolist(slow_interpolator(alphasets)))
    fast_result = np.asarray(pyhf.tensorlib.tolist(fast_interpolator(alphasets)))

    assert slow_result.shape == fast_result.shape

    assert (
        pytest.approx(
            slow_result[~np.isnan(slow_result)].ravel().tolist(), abs=abs_tolerance
        )

)
            ]
        ),
    )
    assert result.shape == (1,)

    suggested_pars = [
        [1.0] * 3 + [1, 2, 3, 4, 5],  # 2 pois 5 norm
        [1.0] * 3 + [-1, -2, -3, -4, -5],  # 2 pois 5 norm
        [1.0] * 3 + [-1, -2, 0, 1, 2],  # 2 pois 5 norm
    ]
    constraint = gaussian_constraint_combined(config, batch_size=3)
    result = default_backend.astensor(
        pyhf.tensorlib.tolist(
            constraint.logpdf(
                pyhf.tensorlib.astensor(config.auxdata),
                pyhf.tensorlib.astensor(suggested_pars),
            )
        )
    )
    assert np.all(
        np.isclose(
            result,
            np.sum(
                [
                    [
                        default_backend.normal_logpdf(data, mu, sigma)
                        for data, mu, sigma in zip(
                            [0, 0, 0, 0, 0], [1, 2, 3, 4, 5], [1.5, 2.0, 1.0, 1.0, 1.0]
                        )
                    ],
                    [

'signal': {
                'type': 'lumi',
                'name': 'lumi',
                'data': {'mask': [True, True, True]},
            },
            'background': {
                'type': 'lumi',
                'name': 'lumi',
                'data': {'mask': [True, True, True]},
            },
        },
    }

    hsc = lumi_combined([('lumi', 'lumi')], mc, mega_mods)

    mod = hsc.apply(pyhf.tensorlib.astensor([0.5]))
    shape = pyhf.tensorlib.shape(mod)
    assert shape == (1, 2, 1, 3)

    mod = np.asarray(pyhf.tensorlib.tolist(mod))
    assert np.allclose(mod[0, 0, 0], [0.5, 0.5, 0.5])
    assert np.allclose(mod[0, 1, 0], [0.5, 0.5, 0.5])

    hsc = lumi_combined([('lumi', 'lumi')], mc, mega_mods, batch_size=4)

    mod = hsc.apply(pyhf.tensorlib.astensor([[1.0], [2.0], [3.0], [4.0]]))
    shape = pyhf.tensorlib.shape(mod)
    assert shape == (1, 2, 4, 3)

    mod = np.asarray(pyhf.tensorlib.tolist(mod))
    assert np.allclose(mod[0, 0, 0], [1.0, 1.0, 1.0])
    assert np.allclose(mod[0, 0, 1], [2.0, 2.0, 2.0])

mod = hsc.apply(pyhf.tensorlib.astensor([2.0, 3.0, 4.0]))
    shape = pyhf.tensorlib.shape(mod)
    assert shape == (2, 2, 1, 3)

    mod = np.asarray(pyhf.tensorlib.tolist(mod))
    assert np.allclose(mod[0, 0, 0], [2.0, 1.0, 1.0])
    assert np.allclose(mod[1, 0, 0], [1.0, 3.0, 4.0])

    hsc = shapefactor_combined(
        [('shapefac1', 'shapefactor'), ('shapefac2', 'shapefactor')],
        mc,
        mega_mods,
        batch_size=4,
    )
    mod = hsc.apply(
        pyhf.tensorlib.astensor(
            [[2.0, 3.0, 4.0], [5.0, 6.0, 7.0], [8.0, 9.0, 10.0], [11.0, 12.0, 13.0]]
        )
    )
    shape = pyhf.tensorlib.shape(mod)
    assert shape == (2, 2, 4, 3)

    mod = np.asarray(pyhf.tensorlib.tolist(mod))
    assert np.allclose(mod[0, 0, 0], [2.0, 1.0, 1.0])
    assert np.allclose(mod[0, 0, 1], [5.0, 1.0, 1.0])
    assert np.allclose(mod[0, 0, 2], [8.0, 1.0, 1.0])
    assert np.allclose(mod[0, 0, 3], [11.0, 1.0, 1.0])

    assert np.allclose(mod[1, 0, 0], [1.0, 3.0, 4.0])
    assert np.allclose(mod[1, 0, 1], [1.0, 6.0, 7.0])
    assert np.allclose(mod[1, 0, 2], [1.0, 9.0, 10.0])
    assert np.allclose(mod[1, 0, 3], [1.0, 12.0, 13.0])

def test_paramviewer_simple_batched(backend):
    pars = pyhf.tensorlib.astensor([[1, 2, 3, 4], [5, 6, 7, 8], [9, 10, 11, 12]])

    parshape = pyhf.tensorlib.shape(pars)

    view = ParamViewer(
        parshape,
        {'hello': {'slice': slice(0, 2)}, 'world': {'slice': slice(3, 4)}},
        ['hello', 'world'],
    )
    par_slice = view.get(pars)

    assert isinstance(view.index_selection, list)
    assert all(
        [len(x) == 3 for x in view.index_selection]
    )  # first dimension is batch dim

    assert pyhf.tensorlib.shape(par_slice) == (3, 3)

def test_code1_validation(backend, do_tensorized_calc):
    histogramssets = [[[[0.9], [1.0], [1.1]]]]
    alphasets = pyhf.tensorlib.astensor([[-2, -1, 0, 1, 2]])
    expected = pyhf.tensorlib.astensor(
        [[[[0.9 ** 2], [0.9], [1.0], [1.1], [1.1 ** 2]]]]
    )

    interpolator = pyhf.interpolators.get(1, do_tensorized_calc=do_tensorized_calc)(
        histogramssets, subscribe=False
    )
    result_deltas = interpolator(alphasets)

    # calculate the actual change
    histogramssets = pyhf.tensorlib.astensor(histogramssets)
    allsets_allhistos_noms_repeated = pyhf.tensorlib.einsum(
        'sa,shb->shab',
        pyhf.tensorlib.ones(pyhf.tensorlib.shape(alphasets)),
        histogramssets[:, :, 1],
    )
    results = allsets_allhistos_noms_repeated * result_deltas