Derive Bitwidth from Tensor

aws-neuron / aws-neuron-sdk

Powering AWS purpose-built machine learning chips. Blazing fast and cost effective, natively integrated into PyTorch and TensorFlow and integrated with your favorite AWS services

https://aws.amazon.com/machine-learning/neuron/

Other

450 stars 152 forks source link

Derive Bitwidth from Tensor #1008

Open nandeeka opened 2 weeks ago

nandeeka commented 2 weeks ago

Is there a way to automatically derive the bitwidth given a type? For example, if I have a tensor A, I can get its type with A.dtype, but can I get the bitwidth of the individual elements? In Numpy, we have itemwidth (see here). Thanks!

aws-qieqingy commented 2 weeks ago

Hi Nandeeka! We currently don't have such API publicly available, and we will add it in the next release. Thanks for your feedback.

nandeeka commented 2 weeks ago

Following a comment from @aws-zhehongb, A.dtype.itemsize (following this documentation) works, since NKI datatypes are just Numpy datatypes.