code: def count_non_zero(X: Tensor, axis, keepdims: bool = False) -> Tensor:
      "Count non-zero elements in tensor."
      non_zeros = cast_to_floatx(
          kbackend.not_equal(X, kbackend.constant(0))
      )
      return kbackend.sum(non_zeros, axis=axis, keepdims=keepdims)
  
imports: from __future__ import annotations
  from typing import TypeVar
  import numpy
  import tensorflow
  import tensorflow.keras.backend
  from innvestigate.backend.types import Layer, OptionalList, ShapeTuple, Tensor
label_desc: Use image node_id 10 for calling a global function with example usage: count_non_zero(X, axis, keepdims) and returns: kbackend