How to use the coffea.util.awkward.JaggedArray.fromcounts function in coffea

extractor.add_weight_sets(["testSF2d scalefactors_Tight_Electron tests/samples/testSF2d.histo.root"])
    
    extractor.finalize(reduce_list=['testSF2d'])

    evaluator = extractor.make_evaluator()

    counts, test_eta, test_pt = dummy_jagged_eta_pt()
    
    # test flat eval
    test_out = evaluator["testSF2d"](test_eta, test_pt)

    # print it
    print(evaluator["testSF2d"])
    
    # test structured eval
    test_eta_jagged = awkward.JaggedArray.fromcounts(counts, test_eta)
    test_pt_jagged = awkward.JaggedArray.fromcounts(counts, test_pt)
    test_out_jagged = evaluator["testSF2d"](test_eta_jagged, test_pt_jagged)

    assert (test_out_jagged.counts==counts).all()
    assert (test_out==test_out_jagged.flatten()).all()

    # From make_expected_lookup.py
    expected_output = np.array([
       0.90780139, 0.82748538, 0.86332178, 0.86332178, 0.97981155,
       0.79701495, 0.88245934, 0.82857144, 0.91884059, 0.97466666,
       0.94072163, 1.00775194, 0.82748538, 1.00775194, 0.97203946,
       0.98199672, 0.80655736, 0.90893763, 0.88245934, 0.79701495,
       0.82748538, 0.82857144, 0.91884059, 0.90893763, 0.97520661,
       0.97520661, 0.82748538, 0.91884059, 0.97203946, 0.88245934,
       0.79701495, 0.9458763 , 1.00775194, 0.80655736, 1.00775194,
       1.00775194, 0.98976982, 0.98976982, 0.86332178, 0.94072163,

# ----> clamps can be different from bins
    # ----> if there is more than one binning variable this array is jagged
    # ----> just make it jagged all the time
    clamp_mins = {}
    clamp_maxs = {}
    var_order = []
    offset_col = 2 * nBinnedVars + 1
    offset_name = nBinnedVars + 2
    jagged_counts = np.ones(bins[bin_order[0]].size - 1, dtype=np.int)
    if len(bin_order) > 1:
        jagged_counts = np.maximum(bins[bin_order[1]].counts - 1, 0)  # need counts-1 since we only care about Nbins
    for i in range(nEvalVars):
        var_order.append(layout[i + offset_name])
        if not interpolatedFunc:
            clamp_mins[layout[i + offset_name]] = awkward.JaggedArray.fromcounts(jagged_counts, np.atleast_1d(pars[columns[i + offset_col]]))
            clamp_maxs[layout[i + offset_name]] = awkward.JaggedArray.fromcounts(jagged_counts, np.atleast_1d(pars[columns[i + offset_col + 1]]))
            assert(clamp_mins[layout[i + offset_name]].valid())
            assert(clamp_maxs[layout[i + offset_name]].valid())
            offset_col += 1

    # now get the parameters, which we will look up with the clamped values
    parms = []
    parm_order = []
    offset_col = 2 * nBinnedVars + 1 + int(not interpolatedFunc) * 2 * nEvalVars
    for i in range(nParms):
        jag = awkward.JaggedArray.fromcounts(jagged_counts, pars[columns[i + offset_col]])
        assert(jag.valid())
        parms.append(jag)
        parm_order.append('p%i' % (i))

    wrapped_up = {}
    wrapped_up[(name, 'jme_standard_function')] = (formula,

bins[layout[i + offset_name]] = np.union1d(binMins, binMaxs[-1:])
        else:
            counts = np.zeros(0, dtype=np.int)
            allBins = np.zeros(0, dtype=np.double)
            for binMin in bins[bin_order[0]][:-1]:
                binMins = np.unique(pars[np.where(pars[columns[0]] == binMin)][columns[i + offset_col]])
                binMaxs = np.unique(pars[np.where(pars[columns[0]] == binMin)][columns[i + offset_col + 1]])
                theBins = None
                if np.all(binMins[1:] == binMaxs[:-1]):
                    theBins = np.union1d(binMins, binMaxs)
                else:
                    warnings.warn('binning for file for %s is malformed in variable %s' % (name, layout[i + offset_name]))
                    theBins = np.union1d(binMins, binMaxs[-1:])
                allBins = np.append(allBins, theBins)
                counts = np.append(counts, theBins.size)
            bins[layout[i + offset_name]] = awkward.JaggedArray.fromcounts(counts, allBins)
        bin_order.append(layout[i + offset_name])
        offset_col += 1

    # skip nvars to the variable columns
    # the columns here define clamps for the variables defined in columns[]
    # ----> clamps can be different from bins
    # ----> if there is more than one binning variable this array is jagged
    # ----> just make it jagged all the time
    clamp_mins = {}
    clamp_maxs = {}
    var_order = []
    offset_col = 2 * nBinnedVars + 1
    offset_name = nBinnedVars + 2
    jagged_counts = np.ones(bins[bin_order[0]].size - 1, dtype=np.int)
    if len(bin_order) > 1:
        jagged_counts = np.maximum(bins[bin_order[1]].counts - 1, 0)  # need counts-1 since we only care about Nbins

def _default_fastmatch(first, second, deltaRCut=10000):
    drCut2 = deltaRCut**2
    args = first.eta._argcross(second.eta)
    argsnested = awkward.JaggedArray.fromcounts(first.eta.counts,
                                                awkward.JaggedArray.fromcounts(first.eta.tojagged(second.eta.counts).flatten(),
                                                                               args._content))
    eta0s = first.eta.content[argsnested.content.content.i0]
    eta1s = second.eta.content[argsnested.content.content.i1]
    phi0s = first.phi.content[argsnested.content.content.i0]
    phi1s = second.phi.content[argsnested.content.content.i1]
    offsets_outer = argsnested.offsets
    offsets_inner = argsnested.content.offsets
    detas = np.abs(eta0s - eta1s)
    dphis = (phi0s - phi1s + math.pi) % (2 * math.pi) - math.pi
    passdr = ((detas**2 + dphis**2) < drCut2)
    passdr = awkward.JaggedArray.fromoffsets(offsets_inner, passdr)
    return awkward.JaggedArray.fromoffsets(offsets_outer, passdr.any())

# the columns here define clamps for the variables defined in columns[]
    # ----> clamps can be different from bins
    # ----> if there is more than one binning variable this array is jagged
    # ----> just make it jagged all the time
    clamp_mins = {}
    clamp_maxs = {}
    var_order = []
    offset_col = 2 * nBinnedVars + 1
    offset_name = nBinnedVars + 2
    jagged_counts = np.ones(bins[bin_order[0]].size - 1, dtype=np.int)
    if len(bin_order) > 1:
        jagged_counts = np.maximum(bins[bin_order[1]].counts - 1, 0)  # need counts-1 since we only care about Nbins
    for i in range(nEvalVars):
        var_order.append(layout[i + offset_name])
        if not interpolatedFunc:
            clamp_mins[layout[i + offset_name]] = awkward.JaggedArray.fromcounts(jagged_counts, np.atleast_1d(pars[columns[i + offset_col]]))
            clamp_maxs[layout[i + offset_name]] = awkward.JaggedArray.fromcounts(jagged_counts, np.atleast_1d(pars[columns[i + offset_col + 1]]))
            assert(clamp_mins[layout[i + offset_name]].valid())
            assert(clamp_maxs[layout[i + offset_name]].valid())
            offset_col += 1

    # now get the parameters, which we will look up with the clamped values
    parms = []
    parm_order = []
    offset_col = 2 * nBinnedVars + 1 + int(not interpolatedFunc) * 2 * nEvalVars
    for i in range(nParms):
        jag = awkward.JaggedArray.fromcounts(jagged_counts, pars[columns[i + offset_col]])
        assert(jag.valid())
        parms.append(jag)
        parm_order.append('p%i' % (i))

    wrapped_up = {}