diff --git a/python/mxnet/ndarray/numpy/_op.py b/python/mxnet/ndarray/numpy/_op.py
index f5de082a846c..58ac10eefebc 100644
--- a/python/mxnet/ndarray/numpy/_op.py
+++ b/python/mxnet/ndarray/numpy/_op.py
@@ -5308,9 +5308,8 @@ def nan_to_num(x, copy=True, nan=0.0, posinf=None, neginf=None, **kwargs):
     else:
         raise TypeError('type {} not supported'.format(str(type(x))))
 
-@set_module('mxnet.ndarray.numpy')
-
 
+@set_module('mxnet.ndarray.numpy')
 def diagflat(arr, k=0):
     """
     Create a two-dimensional array with the flattened input as a diagonal.
diff --git a/src/operator/numpy/np_diagflat_op-inl.h b/src/operator/numpy/np_diagflat_op-inl.h
index 93a648f16103..90be875d5567 100644
--- a/src/operator/numpy/np_diagflat_op-inl.h
+++ b/src/operator/numpy/np_diagflat_op-inl.h
@@ -1,3 +1,26 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+/*!
+ * \file np_diagflat_op-inl.h
+ * \brief Function definition of numpy-compatible diff operator
+ */
 #ifndef MXNET_OPERATOR_NUMPY_DIAGFLAT_OP_INL_H_
 #define MXNET_OPERATOR_NUMPY_DIAGFLAT_OP_INL_H_
 
@@ -36,7 +59,7 @@ inline mxnet::TShape NumpyDiagflatShapeImpl(const mxnet::TShape& ishape, const i
 
     if (ishape.ndim() >=2 ){
     	auto s = 1;
-    	for(uint32_t i = 0; i < ishape.ndim(); i++){
+    	for(int i = 0; i < ishape.ndim(); i++){
     		if(ishape[i] >= 2){
     			s = s * ishape[i];
     		}
diff --git a/src/operator/numpy/np_diagflat_op.cc b/src/operator/numpy/np_diagflat_op.cc
index 888dfe61395d..66bd8544e762 100644
--- a/src/operator/numpy/np_diagflat_op.cc
+++ b/src/operator/numpy/np_diagflat_op.cc
@@ -1,3 +1,26 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+/*!
+ * \file np_diagflat_op-inl.cc
+ * \brief CPU Implementation of numpy-compatible dot
+ */
 #include "./np_diagflat_op-inl.h"
 
 namespace mxnet{
diff --git a/src/operator/numpy/np_diagflat_op.cu b/src/operator/numpy/np_diagflat_op.cu
index 106e5b7a6dec..72993e6093b6 100644
--- a/src/operator/numpy/np_diagflat_op.cu
+++ b/src/operator/numpy/np_diagflat_op.cu
@@ -1,3 +1,27 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+/*!
+ * \file np_diagflat_op-inl.cu
+ * \brief GPU implementation of numpy-compatible diff operator
+ */
+
 #include "./np_diagflat_op-inl.h"
 
 namespace mxnet{