How to convert a method to work on any tensor dimension

1

Hi team,

I have defined a method which takes a decimal value and creates a representation of it know as temporal coding.

def getDecimalToUnary(input, bitwidth, encode="TC"):
    """This method accepts an input in decimal format and generates an unary temporal encoded representation of the input.
        The length of the unary bit stream depends on the bitwidth

    Args:
        input (int): The input decimal that needs to be transistion encoded
        bitwidth (int): The bit stream length is 2**bitwidth
        encode (str, optional): _description_. Defaults to 'TC'.  TC represents temporal encoding
    """
    if encode == "TC":
        unary_code = torch.cat(
            (torch.zeros(2**bitwidth - input), torch.ones(input)), 0
        )
        return unary_code
    else:
        print("Other encoding schemes are not implemented")

Rules of generating temporal code of a decimal are:

  1. The generated temporal code length should be equal to 2**bitwidth
  2. The total number of 1’s in the bitstream should be equal to decimal value.
  3. All the 1’s should be together at the end of the bitstream
    For example:
    If input is 5 and bitwidth is 3. Then, the unary bitstream is of length 2**3 = 8bits.
    The temporal representation of 5 is 00011111.
    The above code is able to work on single decimal value, however I want to have this method work for any dimension tensor.
    For instance, for tensor([1,2],[2,5]) with bitwidth 3, the output of the method should be
    tensor([[[0,0,0,0,0,0,0,1],[0,0,0,0,0,0,1,1]],[[0,0,0,0,0,0,1,1],[0,0,0,1,1,1,1,1]]])