def METHOD_NAME(cls, op: "LearnCountNonzero"):
    input_tensor = op.inputs[0]
    out_tensor = op.outputs[0]
    if op.sample_weight is not None:
        if has_unknown_shape(input_tensor):
            yield
        sample_weight = yield from recursive_tile(
            op.sample_weight.rechunk({0: input_tensor.nsplits[0]})
        )
    else:
        sample_weight = None
    chunks = []
    for input_chunk in input_tensor.chunks:
        if sample_weight is None:
            weight_chunk = None
        else:
            weight_chunk = sample_weight.cix[(input_chunk.index[0],)]
        new_op = op.copy().reset_key()
        new_op.sample_weight = weight_chunk
        inputs = [input_chunk] if not weight_chunk else [input_chunk, weight_chunk]
        if op.axis is None:
            shape = (1, 1)
        elif op.axis == 0:
            shape = (1, input_chunk.shape[1])
        else:
            shape = (input_chunk.shape[0], 1)
        chunks.append(
            new_op.new_chunk(
                inputs, shape=shape, dtype=out_tensor.dtype, index=input_chunk.index
            )
        )
    new_op = op.copy().reset_key()
    if op.axis is None:
        nsplits = tuple((1,) * len(split) for split in input_tensor.nsplits)
        shape = tuple(len(split) for split in input_tensor.nsplits)
    elif op.axis == 0:
        nsplits = ((1,) * len(input_tensor.nsplits[0]), input_tensor.nsplits[1])
        shape = (len(input_tensor.nsplits[0]), input_tensor.shape[1])
    else:
        nsplits = (input_tensor.nsplits[0], (1,) * len(input_tensor.nsplits[1]))
        shape = (input_tensor.shape[0], len(input_tensor.nsplits[1]))
    tileable = new_op.new_tileable(
        out_tensor.inputs,
        chunks=chunks,
        nsplits=nsplits,
        shape=shape,
        dtype=out_tensor.dtype,
    )
    return [(yield from recursive_tile(mt.sum(tileable, axis=op.axis)))]