Too verbose? Or easy to read?

It is easy to read but at the point this is written why not just do:

yield SampleInput(make_arg(M, M), args=(make_bool_mask(M, M), make_arg(M, M)))

instead of each case?

cases = (((M, M), (M, M), (M, M), False),
             ((M, 1, M), (M, M), (M, M, 1), True),
             ((), (), (), False),
             ((M, 1, M), (), (M, M, 1), True),
             ((), (M, M), (), True),)

def generator():
    for shape, mask_shape, other_shape, broadcasts_input in cases:
            yield SampleInput(make_arg(shape),
                              args=(make_bool_mask(mask_shape), make_arg(other_shape)),
                              broadcasts_input=broadcasts_input)

Is it ok to stick to usual cases as tuple. Even yield for every sample input is very verbose and hard to extend the cases if need (ever) be.

This is going to need a comment; this is because we don't support the input being anything but the first argument, right?

Yes. For this specific case we can pass input as condition but since condition is always bool, so we won't gradcheck on it.

Can we use gradcheck_wrapper to handle the gradcheck case? Or are there other tests which would still be an issue?

Also affect test_variant_consistency_eager

OK

Skip will need a comment, too

The function this is replacing guaranteed that at least one value in the mask is True

This PR should preserve that behavior. It can do so by creating its own make_bool_mask function that checks if the tensor contains at least one True value and, if not, inserts a True a "True" at a random point, for example.

Is it okay to use mask_not_all_zero? It is already used by other OpInfo sample functions.

I wonder how this works for index_put?

I believe that args should be on the same device. But for mask neither torch.zeros specifies dtypes not mask_not_all_zero 🤔

Sure; although eventually it'd be good to improve upon it and not just redraw

@mruberry take a look at the updated comment. I am confused.

I think some of the indexing ops do allow the indices (and maybe masks) to be on the CPU:

import torch

t = torch.tensor((1, 2, 3), device='cuda')
indices = torch.tensor((1,))
values = torch.tensor((4,), device='cuda')
t.index_put_((indices,), values)
: tensor([1, 4, 3], device='cuda:0')

How does this look?

    def make_bool_mask(shape):
        # Make sure atleast one element is nonzero,
        # except for empty tensor
        mask_t = make_tensor(shape, dtype=torch.bool, device=device, requires_grad=False)

        if mask_t.sum() == 0:
            if mask_t.numel() == 1 and len(mask_t.shape) > 0:
                mask_t[0] = True
                return mask_t
            elif reduce(lambda x, y: x*y, mask_t.shape, 1) == 0:
                return mask_t

            def random_index(shape):
                return tuple(map(lambda max_idx: random.randint(0, max_idx), shape))

            mask_t[random_index(mask_t.shape)] = True
            return mask_t
        
        return mask_t

The handling of tensors with zero elements looks good but I think that check can be simplified to mask_t having numel == 0? The handling of the general case also looks correct.

I think the tricky case, then, is if shape is () or (1,). For these I think you need two checks, one for numel() == 1 and a then a second check for len(shape) == 0. In both these cases the initial construction of the tensor is irrelevant, because you know the tensor you want to return, so I would put these checks before the make_tensor case, short-circuit on 0 numel, and then only guard the random_index case behind the call to sum()

Tried to incorporate the suggestion
How about?

    def make_bool_mask(shape):
        # Make sure atleast one element is nonzero,
        # except for empty tensor
        mask_t = make_tensor(shape, dtype=torch.bool, device=device, requires_grad=False)

        if mask_t.numel() == 0:
            return mask_t
        elif mask_t.numel() == 1:
            # Added these after make tensor so that
            # we have to specify `dtype`, `device`, etc only
            # once to make_tensor
            mask_t.fill_(True)
            return mask_t

        if mask_t.sum() == 0:
            def random_index(shape):
                return tuple(map(lambda max_idx: random.randint(0, max_idx), shape))

            mask_t[random_index(mask_t.shape)] = True
            return mask_t

        return mask_t

LGTM!