Skip to content
Snippets Groups Projects
Commit cee44f6d authored by Bhatu's avatar Bhatu
Browse files

Support for reduced mean

Adds support the reduce_mean operation in tensorflow.
Consider the example:
  For inputs:
    Tensor of shape(s0,s1,s2,s3)
    reduction axes = [0,3]

  We generate the following program:
    If keep_dim == true
      output is of shape(1,s1,s2,1)
    else
      output is of shape(s1,s2)

    for i1=[0:s1]
      for i2=[0:s2]
        sum = 0
        for i0=[0:s0]
          for i3=[0:s3]
            sum  = sum + input[i0][i1][i2][i3]
        output[i1][i2] = sum / (s0 * s3)        // keep_dim=false
  OR
        output[0][i1][i2][0] = sum / (s0 * s3)  // keep_dim=true

TODO: Also add support for reduced sum.
parent f416af1e
No related branches found
No related tags found
No related merge requests found
Loading
0% Loading or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment