[Question] Creating lazy matrix from a function

[PeterMitrano]

first off thanks for the great library. I'm wondering if there's a better way to create a matrix from a function that only knows how to load a single example. I thought this would work:

from lazyarray import larray
import numpy as np

pretend_this_is_a_file = [
        [1,2,3],
        [1,2,3],
        [1,2,3],
]

def load(i, j):
    print(type(i), type(j))
    return pretend_this_is_a_file[i][j]

So load only works if i and j are integers. Not, for example, lists or arrays of integers. That means this usage fails:

z = larray(load, shape=[3,3])
print(z[0,0])
print(z[0])
print(z[:, 0])

Is there a better way? I have a workaround, which implements the case for when i or j and arrays of integers, but that seems like it should be done by lazyarray

EX:

    def _load_examples(i, j):
        if isinstance(i, np.ndarray):
            results = []
            for _i in i:
                _name = f'{_i}-{j}.pkl.gz'
                _results_filename = results_dir / _name
                if _results_filename.exists():
                    _result = load_gzipped_pickle(_results_filename)[k]
                else:
                    _result = None
                results.append(_result)
            return np.array(results)
        elif isinstance(j, np.ndarray):
            results = []
            for _j in j:
                _name = f'{i}-{_j}.pkl.gz'
                _results_filename = results_dir / _name
                if _results_filename.exists():
                    _result = load_gzipped_pickle(_results_filename)[k]
                else:
                    _result = None
                results.append(_result)
            return np.array(results)
        else:
            name = f'{i}-{j}.pkl.gz'
            results_filename = results_dir / name
            result = load_gzipped_pickle(results_filename)[k]
            return result

If this is a good idea and you're interested in contributions, I could probably help put a PR together