code: def softmax(x):
      """Simple numpy softmax implementation"""
      # (n_samples, n_classes)
      if len(x.shape) == 2:
          row_max = np.max(x, axis=1)
          x -= row_max.reshape((x.shape[0], 1))
          x = np.exp(x)
          row_sum = np.sum(x, axis=1)
          x /= row_sum.reshape((x.shape[0], 1))
      # (n_samples, n_tasks, n_classes)
      elif len(x.shape) == 3:
          row_max = np.max(x, axis=2)
          x -= row_max.reshape(x.shape[:2] + (1,))
          x = np.exp(x)
          row_sum = np.sum(x, axis=2)
          x /= row_sum.reshape(x.shape[:2] + (1,))
      return x
  
imports: from __future__ import print_function
  from __future__ import division
  from __future__ import unicode_literals
  import collections
  import pickle
  import os
  import time
  import warnings
  import numpy
  import pandas
  import tensorflow
  import tempfile
  from deepchem.models import Model
  from deepchem.metrics import from_one_hot
  from deepchem.nn import model_ops
  from deepchem.models.tensorflow_models import utils
  from deepchem.trans import undo_transforms
  from deepchem.utils.save import log
  from deepchem.utils.evaluate import Evaluator
  from deepchem.data import pad_features
  from tensorflow.contrib.layers.python.layers import batch_norm
label_desc: Use image node_id 1 for calling a global function with example usage: softmax(x) and returns: x